Skip to content

Commit

Permalink
more minor adjustments to response types
Browse files Browse the repository at this point in the history
  • Loading branch information
miguelgrinberg committed Oct 22, 2024
1 parent d82c98b commit 9777a93
Show file tree
Hide file tree
Showing 6 changed files with 17 additions and 16 deletions.
9 changes: 8 additions & 1 deletion elasticsearch_dsl/response/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]):
Expand Down
7 changes: 4 additions & 3 deletions examples/async/composite_agg.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand All @@ -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
Expand All @@ -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:
Expand Down
4 changes: 0 additions & 4 deletions examples/async/parent_child.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,6 @@

from elasticsearch_dsl import (
AsyncDocument,
AsyncIndex,
AsyncSearch,
Date,
InnerDoc,
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down
7 changes: 4 additions & 3 deletions examples/composite_agg.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand All @@ -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
Expand All @@ -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:
Expand Down
4 changes: 0 additions & 4 deletions examples/parent_child.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,6 @@
from elasticsearch_dsl import (
Date,
Document,
Index,
InnerDoc,
Join,
Keyword,
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion utils/templates/response.__init__.py.tpl
Original file line number Diff line number Diff line change
Expand Up @@ -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]):
Expand Down

0 comments on commit 9777a93

Please sign in to comment.