From e751ab55c8a61300124fa4ac509a2f8fefc51dcc Mon Sep 17 00:00:00 2001 From: Stefan Wojcik Date: Sun, 7 May 2017 23:28:35 -0400 Subject: [PATCH] cleaner as_pymongo + drop coerce_types --- mongoengine/queryset/base.py | 87 +++++++++++------------------------- tests/queryset/queryset.py | 49 +++++++++++++++----- 2 files changed, 63 insertions(+), 73 deletions(-) diff --git a/mongoengine/queryset/base.py b/mongoengine/queryset/base.py index 7c60a489..a3f41c3c 100644 --- a/mongoengine/queryset/base.py +++ b/mongoengine/queryset/base.py @@ -67,7 +67,6 @@ class BaseQuerySet(object): self._scalar = [] self._none = False self._as_pymongo = False - self._as_pymongo_coerce = False self._search_text = None # If inheritance is allowed, only return instances and instances of @@ -728,11 +727,12 @@ class BaseQuerySet(object): '%s is not a subclass of BaseQuerySet' % new_qs.__name__) copy_props = ('_mongo_query', '_initial_query', '_none', '_query_obj', - '_where_clause', '_loaded_fields', '_ordering', '_snapshot', - '_timeout', '_class_check', '_slave_okay', '_read_preference', - '_iter', '_scalar', '_as_pymongo', '_as_pymongo_coerce', + '_where_clause', '_loaded_fields', '_ordering', + '_snapshot', '_timeout', '_class_check', '_slave_okay', + '_read_preference', '_iter', '_scalar', '_as_pymongo', '_limit', '_skip', '_hint', '_auto_dereference', - '_search_text', 'only_fields', '_max_time_ms', '_comment') + '_search_text', 'only_fields', '_max_time_ms', + '_comment') for prop in copy_props: val = getattr(self, prop) @@ -939,7 +939,8 @@ class BaseQuerySet(object): posts = BlogPost.objects(...).fields(slice__comments=5) - :param kwargs: A set keywors arguments identifying what to include. + :param kwargs: A set of keyword arguments identifying what to + include, exclude, or slice. .. versionadded:: 0.5 """ @@ -1128,16 +1129,12 @@ class BaseQuerySet(object): """An alias for scalar""" return self.scalar(*fields) - def as_pymongo(self, coerce_types=False): + def as_pymongo(self): """Instead of returning Document instances, return raw values from pymongo. - - :param coerce_types: Field types (if applicable) would be use to - coerce types. """ queryset = self.clone() queryset._as_pymongo = True - queryset._as_pymongo_coerce = coerce_types return queryset def max_time_ms(self, ms): @@ -1799,59 +1796,25 @@ class BaseQuerySet(object): return tuple(data) - def _get_as_pymongo(self, row): - # Extract which fields paths we should follow if .fields(...) was - # used. If not, handle all fields. - if not getattr(self, '__as_pymongo_fields', None): - self.__as_pymongo_fields = [] + def _get_as_pymongo(self, doc): + """Clean up a PyMongo doc, removing fields that were only fetched + for the sake of MongoEngine's implementation, and return it. + """ + # Always remove _cls as a MongoEngine's implementation detail. + if '_cls' in doc: + del doc['_cls'] - for field in self._loaded_fields.fields - set(['_cls']): - self.__as_pymongo_fields.append(field) - while '.' in field: - field, _ = field.rsplit('.', 1) - self.__as_pymongo_fields.append(field) + # If the _id was not included in a .only or was excluded in a .exclude, + # remove it from the doc (we always fetch it so that we can properly + # construct documents). + fields = self._loaded_fields + if fields and '_id' in doc and ( + (fields.value == QueryFieldList.ONLY and '_id' not in fields.fields) or + (fields.value == QueryFieldList.EXCLUDE and '_id' in fields.fields) + ): + del doc['_id'] - all_fields = not self.__as_pymongo_fields - - def clean(data, path=None): - path = path or '' - - if isinstance(data, dict): - new_data = {} - for key, value in data.iteritems(): - new_path = '%s.%s' % (path, key) if path else key - - if all_fields: - include_field = True - elif self._loaded_fields.value == QueryFieldList.ONLY: - include_field = new_path in self.__as_pymongo_fields - else: - include_field = new_path not in self.__as_pymongo_fields - - if include_field: - new_data[key] = clean(value, path=new_path) - data = new_data - elif isinstance(data, list): - data = [clean(d, path=path) for d in data] - else: - if self._as_pymongo_coerce: - # If we need to coerce types, we need to determine the - # type of this field and use the corresponding - # .to_python(...) - EmbeddedDocumentField = _import_class('EmbeddedDocumentField') - - obj = self._document - for chunk in path.split('.'): - obj = getattr(obj, chunk, None) - if obj is None: - break - elif isinstance(obj, EmbeddedDocumentField): - obj = obj.document_type - if obj and data is not None: - data = obj.to_python(data) - return data - - return clean(row) + return doc def _sub_js_fields(self, code): """When fields are specified with [~fieldname] syntax, where diff --git a/tests/queryset/queryset.py b/tests/queryset/queryset.py index 7b12625c..7344fbd9 100644 --- a/tests/queryset/queryset.py +++ b/tests/queryset/queryset.py @@ -4392,21 +4392,44 @@ class QuerySetTest(unittest.TestCase): self.assertEqual(doc_objects, Doc.objects.from_json(json_data)) def test_as_pymongo(self): - from decimal import Decimal + class LastLogin(EmbeddedDocument): + location = StringField() + ip = StringField() + class User(Document): id = ObjectIdField('_id') name = StringField() age = IntField() price = DecimalField() + last_login = EmbeddedDocumentField(LastLogin) User.drop_collection() - User(name="Bob Dole", age=89, price=Decimal('1.11')).save() - User(name="Barack Obama", age=51, price=Decimal('2.22')).save() + + User.objects.create(name="Bob Dole", age=89, price=Decimal('1.11')) + User.objects.create( + name="Barack Obama", + age=51, + price=Decimal('2.22'), + last_login=LastLogin( + location='White House', + ip='104.107.108.116' + ) + ) + + results = User.objects.as_pymongo() + self.assertEqual( + set(results[0].keys()), + set(['_id', 'name', 'age', 'price']) + ) + self.assertEqual( + set(results[1].keys()), + set(['_id', 'name', 'age', 'price', 'last_login']) + ) results = User.objects.only('id', 'name').as_pymongo() - self.assertEqual(sorted(results[0].keys()), sorted(['_id', 'name'])) + self.assertEqual(set(results[0].keys()), set(['_id', 'name'])) users = User.objects.only('name', 'price').as_pymongo() results = list(users) @@ -4417,16 +4440,20 @@ class QuerySetTest(unittest.TestCase): self.assertEqual(results[1]['name'], 'Barack Obama') self.assertEqual(results[1]['price'], 2.22) - # Test coerce_types - users = User.objects.only( - 'name', 'price').as_pymongo(coerce_types=True) + users = User.objects.only('name', 'last_login').as_pymongo() results = list(users) self.assertTrue(isinstance(results[0], dict)) self.assertTrue(isinstance(results[1], dict)) - self.assertEqual(results[0]['name'], 'Bob Dole') - self.assertEqual(results[0]['price'], Decimal('1.11')) - self.assertEqual(results[1]['name'], 'Barack Obama') - self.assertEqual(results[1]['price'], Decimal('2.22')) + self.assertEqual(results[0], { + 'name': 'Bob Dole' + }) + self.assertEqual(results[1], { + 'name': 'Barack Obama', + 'last_login': { + 'location': 'White House', + 'ip': '104.107.108.116' + } + }) def test_as_pymongo_json_limit_fields(self):