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