more minor adjustments to response types

Author: miguelgrinberg

elasticsearch_dsl/response/__init__.py

@@ -40,7 +40,14 @@
     from ..search_base import Request, SearchBase
     from ..update_by_query_base import UpdateByQueryBase
 
-__all__ = ["Response", "AggResponse", "UpdateByQueryResponse", "Hit", "HitMeta"]
+__all__ = [
+    "Response",
+    "AggResponse",
+    "UpdateByQueryResponse",
+    "Hit",
+    "HitMeta",
+    "AggregateResponseType",
+]
 
 
 class Response(AttrDict[Any], Generic[_R]):

examples/async/composite_agg.py

@@ -17,11 +17,12 @@
 
 import asyncio
 import os
-from typing import Any, AsyncIterator, Dict, Mapping, Sequence
+from typing import Any, AsyncIterator, Dict, Mapping, Sequence, cast
 
 from elasticsearch.helpers import async_bulk
 
 from elasticsearch_dsl import Agg, AsyncSearch, Response, aggs, async_connections
+from elasticsearch_dsl.types import CompositeAggregate
 from tests.test_integration.test_data import DATA, GIT_INDEX
 
 
@@ -30,7 +31,7 @@ async def scan_aggs(
     source_aggs: Sequence[Mapping[str, Agg]],
     inner_aggs: Dict[str, Agg] = {},
     size: int = 10,
-) -> AsyncIterator[Any]:
+) -> AsyncIterator[CompositeAggregate]:
     """
     Helper function used to iterate over all possible bucket combinations of
     ``source_aggs``, returning results of ``inner_aggs`` for each. Uses the
@@ -54,7 +55,7 @@ async def run_search(**kwargs: Any) -> Response:
     response = await run_search()
     while response.aggregations["comp"].buckets:
         for b in response.aggregations["comp"].buckets:
-            yield b
+            yield cast(CompositeAggregate, b)
         if "after_key" in response.aggregations["comp"]:
             after = response.aggregations["comp"].after_key
         else:

examples/async/parent_child.py

@@ -46,7 +46,6 @@
 
 from elasticsearch_dsl import (
     AsyncDocument,
-    AsyncIndex,
     AsyncSearch,
     Date,
     InnerDoc,
@@ -92,7 +91,7 @@ class Post(AsyncDocument):
         # definitions here help type checkers understand additional arguments
         # that are allowed in the constructor
         _routing: str = mapped_field(default=None)
-        _index: AsyncIndex = mapped_field(default=None)
+        _index: str = mapped_field(default=None)
         _id: Optional[int] = mapped_field(default=None)
 
     created: Optional[datetime] = mapped_field(default=None)

examples/composite_agg.py

@@ -16,11 +16,12 @@
 #  under the License.
 
 import os
-from typing import Any, Dict, Iterator, Mapping, Sequence
+from typing import Any, Dict, Iterator, Mapping, Sequence, cast
 
 from elasticsearch.helpers import bulk
 
 from elasticsearch_dsl import Agg, Response, Search, aggs, connections
+from elasticsearch_dsl.types import CompositeAggregate
 from tests.test_integration.test_data import DATA, GIT_INDEX
 
 
@@ -29,7 +30,7 @@ def scan_aggs(
     source_aggs: Sequence[Mapping[str, Agg]],
     inner_aggs: Dict[str, Agg] = {},
     size: int = 10,
-) -> Iterator[Any]:
+) -> Iterator[CompositeAggregate]:
     """
     Helper function used to iterate over all possible bucket combinations of
     ``source_aggs``, returning results of ``inner_aggs`` for each. Uses the
@@ -53,7 +54,7 @@ def run_search(**kwargs: Any) -> Response:
     response = run_search()
     while response.aggregations["comp"].buckets:
         for b in response.aggregations["comp"].buckets:
-            yield b
+            yield cast(CompositeAggregate, b)
         if "after_key" in response.aggregations["comp"]:
             after = response.aggregations["comp"].after_key
         else:

examples/parent_child.py

@@ -46,7 +46,6 @@
 from elasticsearch_dsl import (
     Date,
     Document,
-    Index,
     InnerDoc,
     Join,
     Keyword,
@@ -91,7 +90,7 @@ class Post(Document):
         # definitions here help type checkers understand additional arguments
         # that are allowed in the constructor
         _routing: str = mapped_field(default=None)
-        _index: Index = mapped_field(default=None)
+        _index: str = mapped_field(default=None)
         _id: Optional[int] = mapped_field(default=None)
 
     created: Optional[datetime] = mapped_field(default=None)

utils/templates/response.__init__.py.tpl

@@ -40,7 +40,7 @@ if TYPE_CHECKING:
     from ..update_by_query_base import UpdateByQueryBase
     from .. import types
 
-__all__ = ["Response", "AggResponse", "UpdateByQueryResponse", "Hit", "HitMeta"]
+__all__ = ["Response", "AggResponse", "UpdateByQueryResponse", "Hit", "HitMeta", "AggregateResponseType"]
 
 
 class Response(AttrDict[Any], Generic[_R]):