1# Licensed to Elasticsearch B.V. under one or more contributor 2# license agreements. See the NOTICE file distributed with 3# this work for additional information regarding copyright 4# ownership. Elasticsearch B.V. licenses this file to you under 5# the Apache License, Version 2.0 (the "License"); you may 6# not use this file except in compliance with the License. 7# You may obtain a copy of the License at 8# 9# http://www.apache.org/licenses/LICENSE-2.0 10# 11# Unless required by applicable law or agreed to in writing, 12# software distributed under the License is distributed on an 13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14# KIND, either express or implied. See the License for the 15# specific language governing permissions and limitations 16# under the License. 17 18from ..utils import AttrDict, AttrList 19from . import Response, AggResponse 20 21 22class Bucket(AggResponse): 23 def __init__(self, aggs, search, data, field=None): 24 super(Bucket, self).__init__(aggs, search, data) 25 26 27class FieldBucket(Bucket): 28 def __init__(self, aggs, search, data, field=None): 29 if field: 30 data["key"] = field.deserialize(data["key"]) 31 super(FieldBucket, self).__init__(aggs, search, data, field) 32 33 34class BucketData(AggResponse): 35 _bucket_class = Bucket 36 37 def _wrap_bucket(self, data): 38 return self._bucket_class( 39 self._meta["aggs"], 40 self._meta["search"], 41 data, 42 field=self._meta.get("field"), 43 ) 44 45 def __iter__(self): 46 return iter(self.buckets) 47 48 def __len__(self): 49 return len(self.buckets) 50 51 def __getitem__(self, key): 52 if isinstance(key, (int, slice)): 53 return self.buckets[key] 54 return super(BucketData, self).__getitem__(key) 55 56 @property 57 def buckets(self): 58 if not hasattr(self, "_buckets"): 59 field = getattr(self._meta["aggs"], "field", None) 60 if field: 61 self._meta["field"] = self._meta["search"]._resolve_field(field) 62 bs = self._d_["buckets"] 63 if isinstance(bs, list): 64 bs = AttrList(bs, obj_wrapper=self._wrap_bucket) 65 else: 66 bs = AttrDict({k: self._wrap_bucket(bs[k]) for k in bs}) 67 super(AttrDict, self).__setattr__("_buckets", bs) 68 return self._buckets 69 70 71class FieldBucketData(BucketData): 72 _bucket_class = FieldBucket 73 74 75class TopHitsData(Response): 76 def __init__(self, agg, search, data): 77 super(AttrDict, self).__setattr__( 78 "meta", AttrDict({"agg": agg, "search": search}) 79 ) 80 super(TopHitsData, self).__init__(search, data) 81