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 099a52f commit b6f3fe4
Show file tree
Hide file tree
Showing 6 changed files with 19 additions and 12 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
3 changes: 1 addition & 2 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,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)
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
3 changes: 1 addition & 2 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,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)
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 b6f3fe4

Please sign in to comment.