Skip to content

Commit

Permalink
buckets_as_dict property
Browse files Browse the repository at this point in the history
  • Loading branch information
miguelgrinberg committed Oct 18, 2024
1 parent 1caae1d commit 18e596e
Show file tree
Hide file tree
Showing 5 changed files with 201 additions and 57 deletions.
15 changes: 7 additions & 8 deletions elasticsearch_dsl/response/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,6 @@
Optional,
Sequence,
Tuple,
TypedDict,
Union,
cast,
)
Expand Down Expand Up @@ -196,7 +195,7 @@ def search_after(self) -> "SearchBase[_R]":
return self._search.extra(search_after=self.hits[-1].meta.sort) # type: ignore


_Aggregate = Union[
AggregateResponseType = Union[
"types.CardinalityAggregate",
"types.HdrPercentilesAggregate",
"types.HdrPercentileRanksAggregate",
Expand Down Expand Up @@ -268,28 +267,28 @@ def search_after(self) -> "SearchBase[_R]":
"types.MatrixStatsAggregate",
"types.GeoLineAggregate",
]
_AggResponseMeta = TypedDict(
"_AggResponseMeta", {"search": "Request[_R]", "aggs": Mapping[str, _Aggregate]}
)


class AggResponse(AttrDict[Any], Generic[_R]):
"""An Elasticsearch aggregation response."""

_meta: Dict[str, Any]

def __init__(self, aggs: "Agg[_R]", search: "Request[_R]", data: Dict[str, Any]):
super(AttrDict, self).__setattr__("_meta", {"search": search, "aggs": aggs})
super().__init__(data)

def __getitem__(self, attr_name: str) -> _Aggregate:
def __getitem__(self, attr_name: str) -> AggregateResponseType:
if attr_name in self._meta["aggs"]:
# don't do self._meta['aggs'][attr_name] to avoid copying
agg = self._meta["aggs"].aggs[attr_name]
return cast(
_Aggregate, agg.result(self._meta["search"], self._d_[attr_name])
AggregateResponseType,
agg.result(self._meta["search"], self._d_[attr_name]),
)
return super().__getitem__(attr_name) # type: ignore

def __iter__(self) -> Iterator[_Aggregate]: # type: ignore[override]
def __iter__(self) -> Iterator[AggregateResponseType]: # type: ignore[override]
for name in self._meta["aggs"]:
yield self[name]

Expand Down
Loading

0 comments on commit 18e596e

Please sign in to comment.