diff --git a/elasticsearch_dsl/response/__init__.py b/elasticsearch_dsl/response/__init__.py index fcc5bdc3..eea1b87f 100644 --- a/elasticsearch_dsl/response/__init__.py +++ b/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]): diff --git a/examples/async/composite_agg.py b/examples/async/composite_agg.py index 1e5487aa..37c7bfcb 100644 --- a/examples/async/composite_agg.py +++ b/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: diff --git a/examples/async/parent_child.py b/examples/async/parent_child.py index 80c8f63a..6668a77c 100644 --- a/examples/async/parent_child.py +++ b/examples/async/parent_child.py @@ -46,7 +46,6 @@ from elasticsearch_dsl import ( AsyncDocument, - AsyncIndex, AsyncSearch, Date, InnerDoc, @@ -92,7 +91,6 @@ 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) _id: Optional[int] = mapped_field(default=None) created: Optional[datetime] = mapped_field(default=None) @@ -161,8 +159,6 @@ async def add_answer( answer = Answer( # required make sure the answer is stored in the same shard _routing=self.meta.id, - # since we don't have explicit index, ensure same index as self - _index=self.meta.index, # set up the parent/child mapping question_answer={"name": "answer", "parent": self.meta.id}, # pass in the field values diff --git a/examples/composite_agg.py b/examples/composite_agg.py index 2367669a..755c0a8c 100644 --- a/examples/composite_agg.py +++ b/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: diff --git a/examples/parent_child.py b/examples/parent_child.py index 25a2e9a7..5acbbd72 100644 --- a/examples/parent_child.py +++ b/examples/parent_child.py @@ -46,7 +46,6 @@ from elasticsearch_dsl import ( Date, Document, - Index, InnerDoc, Join, Keyword, @@ -91,7 +90,6 @@ 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) _id: Optional[int] = mapped_field(default=None) created: Optional[datetime] = mapped_field(default=None) @@ -160,8 +158,6 @@ def add_answer( answer = Answer( # required make sure the answer is stored in the same shard _routing=self.meta.id, - # since we don't have explicit index, ensure same index as self - _index=self.meta.index, # set up the parent/child mapping question_answer={"name": "answer", "parent": self.meta.id}, # pass in the field values diff --git a/utils/templates/response.__init__.py.tpl b/utils/templates/response.__init__.py.tpl index a9aafc01..f9ae5c4e 100644 --- a/utils/templates/response.__init__.py.tpl +++ b/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]):