Generate types for the Elasticsearch responses (#1929)

* generate types for the Elasticsearch response * typing of the update by query response * minor typing updates * Update utils/templates/response.__init__.py.tpl Co-authored-by: Quentin Pradet <[email protected]> --------- Co-authored-by: Quentin Pradet <[email protected]>
elastic · Oct 18, 2024 · 579f572 · 579f572
1 parent e001424
commit 579f572
Show file tree

Hide file tree

Showing 10 changed files with 3,470 additions and 2,368 deletions.
diff --git a/elasticsearch_dsl/document_base.py b/elasticsearch_dsl/document_base.py
@@ -165,7 +165,7 @@ def __init__(self, name: str, bases: Tuple[type, ...], attrs: Dict[str, Any]):
         fields.update(annotations.keys())
         field_defaults = {}
         for name in fields:
-            value = None
+            value: Any = None
             required = None
             multi = None
             if name in annotations:
@@ -201,7 +201,7 @@ def __init__(self, name: str, bases: Tuple[type, ...], attrs: Dict[str, Any]):
                     field_args = [type_]
                 elif type_ in self.type_annotation_map:
                     # use best field type for the type hint provided
-                    field, field_kwargs = self.type_annotation_map[type_]
+                    field, field_kwargs = self.type_annotation_map[type_]  # type: ignore
 
                 if field:
                     field_kwargs = {

diff --git a/elasticsearch_dsl/response/__init__.py b/elasticsearch_dsl/response/__init__.py
@@ -22,7 +22,9 @@
     Generic,
     Iterator,
     List,
+    Mapping,
     Optional,
+    Sequence,
     Tuple,
     Union,
     cast,
@@ -32,6 +34,7 @@
 from .hit import Hit, HitMeta
 
 if TYPE_CHECKING:
+    from .. import types
     from ..aggs import Agg
     from ..faceted_search_base import FacetedSearchBase
     from ..search_base import Request, SearchBase
@@ -41,11 +44,47 @@
 
 
 class Response(AttrDict[Any], Generic[_R]):
+    """An Elasticsearch response.
+
+    :arg took: (required)
+    :arg timed_out: (required)
+    :arg _shards: (required)
+    :arg hits: search results
+    :arg aggregations: aggregation results
+    :arg _clusters:
+    :arg fields:
+    :arg max_score:
+    :arg num_reduce_phases:
+    :arg profile:
+    :arg pit_id:
+    :arg _scroll_id:
+    :arg suggest:
+    :arg terminated_early:
+    """
+
     _search: "SearchBase[_R]"
     _faceted_search: "FacetedSearchBase[_R]"
     _doc_class: Optional[_R]
     _hits: List[_R]
 
+    took: int
+    timed_out: bool
+    _shards: "types.ShardStatistics"
+    _clusters: "types.ClusterStatistics"
+    fields: Mapping[str, Any]
+    max_score: float
+    num_reduce_phases: int
+    profile: "types.Profile"
+    pit_id: str
+    _scroll_id: str
+    suggest: Mapping[
+        str,
+        Sequence[
+            Union["types.CompletionSuggest", "types.PhraseSuggest", "types.TermSuggest"]
+        ],
+    ]
+    terminated_early: bool
+
     def __init__(
         self,
         search: "Request[_R]",
@@ -176,8 +215,45 @@ def __iter__(self) -> Iterator["Agg"]:  # type: ignore[override]
 
 
 class UpdateByQueryResponse(AttrDict[Any], Generic[_R]):
+    """An Elasticsearch update by query response.
+
+    :arg batches:
+    :arg failures:
+    :arg noops:
+    :arg deleted:
+    :arg requests_per_second:
+    :arg retries:
+    :arg task:
+    :arg timed_out:
+    :arg took:
+    :arg total:
+    :arg updated:
+    :arg version_conflicts:
+    :arg throttled:
+    :arg throttled_millis:
+    :arg throttled_until:
+    :arg throttled_until_millis:
+    """
+
     _search: "UpdateByQueryBase[_R]"
 
+    batches: int
+    failures: Sequence["types.BulkIndexByScrollFailure"]
+    noops: int
+    deleted: int
+    requests_per_second: float
+    retries: "types.Retries"
+    task: Union[str, int]
+    timed_out: bool
+    took: Any
+    total: int
+    updated: int
+    version_conflicts: int
+    throttled: Any
+    throttled_millis: Any
+    throttled_until: Any
+    throttled_until_millis: Any
+
     def __init__(
         self,
         search: "Request[_R]",