Don't force aggs to remember their name, make them as other DSLs

Author: honzakral

elasticsearch_dsl/aggs.py

@@ -1,21 +1,17 @@
-from six import add_metaclass, iteritems
+from six import add_metaclass
 
-from .utils import DslMeta, DslBase
+from .utils import DslMeta, DslBase, _make_dsl_class
 
 class AggMeta(DslMeta):
     _classes = {}
 
-def A(name_or_agg, agg_type=None, **params):
-    # {"per_tag": {"terms": {"field": "tags"}, "aggs": {...}}}
+def A(name_or_agg, **params):
+    # {"terms": {"field": "tags"}, "aggs": {...}}
     if isinstance(name_or_agg, dict):
-        if params or agg_type or len(name_or_agg) != 1:
+        if params:
             raise #XXX
-        name, agg = name_or_agg.copy().popitem()
-        # {"per_tag": Terms(...)} - happens when copying buckets
-        if isinstance(agg, Agg):
-            return agg
         # copy to avoid modifying in-place
-        agg = agg.copy()
+        agg = name_or_agg.copy()
         # pop out nested aggs
         aggs = agg.pop('aggs', None)
         # should be {"terms": {"fied": "tags"}}
@@ -25,45 +21,23 @@ def A(name_or_agg, agg_type=None, **params):
         if aggs:
             params = params.copy()
             params['aggs'] = aggs
-        return Agg.get_dsl_class(agg_type)(name, **params)
+        return Agg.get_dsl_class(agg_type)(**params)
 
     # Terms(...) just return the nested agg
     elif isinstance(name_or_agg, Agg):
-        if params or agg_type:
+        if params:
             raise #XXX
         return name_or_agg
 
-    elif agg_type is None:
-        raise #XXX
-    
-    # "per_tag", "terms", field="tags"
-    return Agg.get_dsl_class(agg_type)(name_or_agg, **params)
+    # "terms", field="tags"
+    return Agg.get_dsl_class(name_or_agg)(**params)
 
 @add_metaclass(AggMeta)
 class Agg(DslBase):
     _type_name = 'agg'
     _type_shortcut = staticmethod(A)
     name = None
 
-    def __init__(self, name, **params):
-        self._name = name
-        super(Agg, self).__init__(**params)
-
-    def __repr__(self):
-        return '%s(%r, %r%s)' % (
-            self._type_shortcut.__name__,
-            self._name, self.name, self._repr_params()
-        )
-
-    def to_dict(self):
-        d = super(Agg, self).to_dict()
-        # wrap the dict
-        out = {self._name: d}
-        # pop out the nested aggs param to the same level
-        if 'aggs' in d[self.name]:
-            d['aggs'] = d[self.name].pop('aggs')
-        return out
-
 class AggBase(object):
     _param_defs = {
         'aggs': {'type': 'agg', 'hash': True},
@@ -74,7 +48,7 @@ def __getitem__(self, agg_name):
         # make sure we're not mutating a shared state - whenever accessing a
         # bucket, return a shallow copy of it to be safe
         if isinstance(agg, Bucket):
-            agg = A(agg_name, agg.name, **agg._params)
+            agg = A(agg.name, **agg._params)
             # be sure to store the copy so any modifications to it will affect us
             self._params['aggs'][agg_name] = agg
 
@@ -84,7 +58,7 @@ def __setitem__(self, agg_name, agg):
         self.aggs[agg_name] = A(agg)
 
     def _agg(self, bucket, name, agg_type, **params):
-        agg = self[name] = A(name, agg_type, **params)
+        agg = self[name] = A(agg_type, **params)
 
         # For chaining - when creating new buckets return them...
         if bucket:
@@ -101,11 +75,17 @@ def bucket(self, name, agg_type, **params):
 
 
 class Bucket(AggBase, Agg):
-    def __init__(self, name, **params):
-        super(Bucket, self).__init__(name, **params)
+    def __init__(self, **params):
+        super(Bucket, self).__init__(**params)
         # remember self for chaining
         self._base = self
 
+    def to_dict(self):
+        d = super(AggBase, self).to_dict()
+        if 'aggs' in d[self.name]:
+            d['aggs'] = d[self.name].pop('aggs')
+        return d
+
 class Terms(Bucket):
     name = 'terms'
 

elasticsearch_dsl/search.py

@@ -1,3 +1,5 @@
+from six import iteritems
+
 from .query import Q, EMPTY_QUERY, Filtered
 from .filter import F, EMPTY_FILTER
 from .aggs import A, AggBase
@@ -167,7 +169,7 @@ def update_from_dict(self, d):
         if aggs:
             self.aggs._params = {
                 'aggs': dict(
-                    (name, A({name: value})) for (name, value) in aggs.items())
+                    (name, A(value)) for (name, value) in iteritems(aggs))
             }
         if 'sort' in d:
             self._sort = d.pop('sort')

elasticsearch_dsl/utils.py

@@ -149,11 +149,7 @@ def _setattr(self, name, value):
 
                 # dict(name -> DslBase), make sure we pickup all the objs
                 elif pinfo.get('hash'):
-                    d = {}
-                    for k, v in iteritems(value):
-                        v = shortcut({k: v})
-                        d[v._name] = v
-                    value = d
+                    value = dict((k, shortcut(v)) for (k, v) in iteritems(value))
 
                 # single value object, just convert
                 else:
@@ -206,18 +202,15 @@ def to_dict(self):
 
                 # squash all the hash values into one dict
                 elif pinfo.get('hash'):
-                    new_value = {}
-                    for v in value.values():
-                        new_value.update(v.to_dict())
-                    value = new_value
+                    value = dict((k, v.to_dict()) for k, v in iteritems(value))
 
                 # serialize single values
                 else:
                     value = value.to_dict()
 
             # serialize anything with to_dict method
             elif hasattr(value, 'to_dict'):
-                value - value.to_dict()
+                value = value.to_dict()
 
             d[pname] = value
         return {self.name: d}

test_elasticsearch_dsl/test_aggs.py

@@ -5,142 +5,115 @@
 from pytest import raises
 
 def test_repr():
-    max_score = aggs.Max('max_score', field='score')
-    a = aggs.A('per_tag', 'terms', field='tags', aggs={'max_score': max_score})
+    max_score = aggs.Max(field='score')
+    a = aggs.A('terms', field='tags', aggs={'max_score': max_score})
 
-    assert "A('per_tag', 'terms', aggs={'max_score': A('max_score', 'max', field='score')}, field='tags')" == repr(a)
+    assert "A('terms', aggs={'max_score': A('max', field='score')}, field='tags')" == repr(a)
 
 def test_A_creates_proper_agg():
-    a = aggs.A('per_tag', 'terms', field='tags')
+    a = aggs.A('terms', field='tags')
 
     assert isinstance(a, aggs.Terms)
     assert a._params == {'field': 'tags'}
-    assert a._name == 'per_tag'
 
 def test_A_handles_nested_aggs_properly():
-    max_score = aggs.Max('max_score', field='score')
-    a = aggs.A('per_tag', 'terms', field='tags', aggs={'max_score': max_score})
+    max_score = aggs.Max(field='score')
+    a = aggs.A('terms', field='tags', aggs={'max_score': max_score})
 
     assert isinstance(a, aggs.Terms)
     assert a._params == {'field': 'tags', 'aggs': {'max_score': max_score}}
-    assert a._name == 'per_tag'
 
 def test_A_passes_aggs_through():
-    a = aggs.A('per_tag', 'terms', field='tags')
+    a = aggs.A('terms', field='tags')
     assert aggs.A(a) is a
 
 def test_A_from_dict():
     d = {
-        'per_tag': {
-            'terms': {'field': 'tags'},
-            'aggs': {'per_author': {'terms': {'field': 'author.raw'}}},
-        }
+        'terms': {'field': 'tags'},
+        'aggs': {'per_author': {'terms': {'field': 'author.raw'}}},
     }
     a = aggs.A(d)
 
     assert isinstance(a, aggs.Terms)
-    assert a._params == {'field': 'tags', 'aggs': {'per_author': aggs.A('per_author', 'terms', field='author.raw')}}
-    assert a._name == 'per_tag'
-    assert a['per_author'] == aggs.A('per_author', 'terms', field='author.raw')
-    assert a.aggs.per_author == aggs.A('per_author', 'terms', field='author.raw')
+    assert a._params == {'field': 'tags', 'aggs': {'per_author': aggs.A('terms', field='author.raw')}}
+    assert a['per_author'] == aggs.A('terms', field='author.raw')
+    assert a.aggs.per_author == aggs.A('terms', field='author.raw')
 
 def test_A_fails_with_incorrect_dict():
     correct_d = {
-        'per_tag': {
-            'terms': {'field': 'tags'},
-            'aggs': {'per_author': {'terms': {'field': 'author.raw'}}},
-        }
+        'terms': {'field': 'tags'},
+        'aggs': {'per_author': {'terms': {'field': 'author.raw'}}},
     }
 
     with raises(Exception):
         aggs.A(correct_d, field='f')
 
-    with raises(Exception):
-        aggs.A(correct_d, 'name')
-
-    d = deepcopy(correct_d)
-    del d['per_tag']['terms']
-    with raises(Exception):
-        aggs.A(d)
-
-    d = deepcopy(correct_d)
-    d['per_tag']['xx'] = {}
+    d = correct_d.copy()
+    del d['terms']
     with raises(Exception):
         aggs.A(d)
 
-    d = deepcopy(correct_d)
+    d = correct_d.copy()
     d['xx'] = {}
     with raises(Exception):
         aggs.A(d)
 
-def test_A_fails_without_agg_type():
-    with raises(Exception):
-        aggs.A('name', field='f')
-
-def test_A_fails_with_agg_and_name_or_params():
-    a = aggs.A('per_tag', 'terms', field='tags')
-
-    with raises(Exception):
-        aggs.A(a, 'name')
+def test_A_fails_with_agg_and_params():
+    a = aggs.A('terms', field='tags')
 
     with raises(Exception):
         aggs.A(a, field='score')
 
 def test_buckets_are_nestable():
-    a = aggs.Terms('per_tag', field='tags')
+    a = aggs.Terms(field='tags')
     b = a.bucket('per_author', 'terms', field='author.raw')
 
     assert isinstance(b, aggs.Terms)
     assert b._params == {'field': 'author.raw'}
-    assert b._name == 'per_author'
     assert a.aggs == {'per_author': b}
 
 def test_metric_inside_buckets():
-    a = aggs.Terms('per_tag', field='tags')
+    a = aggs.Terms(field='tags')
     b = a.metric('max_score', 'max', field='score')
 
     # returns bucket so it's chainable
     assert a is b
-    assert a.aggs['max_score'] == aggs.Max('max_score', field='score')
+    assert a.aggs['max_score'] == aggs.Max(field='score')
 
 def test_buckets_equals_counts_subaggs():
-    a = aggs.Terms('per_tag', field='tags')
+    a = aggs.Terms(field='tags')
     a.bucket('per_author', 'terms', field='author.raw')
-    b = aggs.Terms('per_tag', field='tags')
+    b = aggs.Terms(field='tags')
 
     assert a != b
 
 def test_buckets_to_dict():
-    a = aggs.Terms('per_tag', field='tags')
+    a = aggs.Terms(field='tags')
     a.bucket('per_author', 'terms', field='author.raw')
 
     assert {
-        'per_tag': {
-            'terms': {'field': 'tags'},
-            'aggs': {'per_author': {'terms': {'field': 'author.raw'}}},
-        }
+        'terms': {'field': 'tags'},
+        'aggs': {'per_author': {'terms': {'field': 'author.raw'}}},
     } == a.to_dict()
 
-    a = aggs.Terms('per_tag', field='tags')
+    a = aggs.Terms(field='tags')
     a.metric('max_score', 'max', field='score')
 
     assert {
-        'per_tag': {
-            'terms': {'field': 'tags'},
-            'aggs': {'max_score': {'max': {'field': 'score'}}},
-        }
+        'terms': {'field': 'tags'},
+        'aggs': {'max_score': {'max': {'field': 'score'}}},
     } == a.to_dict()
 
 def test_nested_buckets_are_reachable_as_getitem():
-    a = aggs.Terms('per_tag', field='tags')
+    a = aggs.Terms(field='tags')
     b = a.bucket('per_author', 'terms', field='author.raw')
 
     assert a['per_author'] is not b
     assert a['per_author'] == b
 
 def test_nested_buckets_are_settable_as_getitem():
-    a = aggs.Terms('per_tag', field='tags')
-    b = a['per_author'] = aggs.A('per_author', 'terms', field='author.raw')
+    a = aggs.Terms(field='tags')
+    b = a['per_author'] = aggs.A('terms', field='author.raw')
 
     assert a.aggs['per_author'] is b
 

test_elasticsearch_dsl/test_search.py

@@ -1,3 +1,4 @@
+from copy import deepcopy
 from mock import Mock
 
 from elasticsearch_dsl import search, query, F, Q
@@ -239,8 +240,12 @@ def test_reverse():
         "size": 5
     }
 
+    d2 = deepcopy(d)
+
     s = search.Search.from_dict(d)
 
+    # make sure we haven't modified anything in place
+    assert d == d2
     assert {"size": 5} == s._extra
     assert d == s.to_dict()