Fix conflicts

Author: RStatkevych

mongoengine/base/document.py

@@ -300,7 +300,7 @@ def get_text_score(self):
 
         return self._data['_text_score']
 
-    def to_mongo(self, use_db_field=True, fields=None):
+    def to_mongo(self, use_db_field=True, fields=None, populate=False):
         """
         Return as SON data ready for use with MongoDB.
         """
@@ -310,11 +310,21 @@ def to_mongo(self, use_db_field=True, fields=None):
         data = SON()
         data["_id"] = None
         data['_cls'] = self._class_name
-        
+
+        Document = _import_class("Document")
+        EmbeddedDocumentField = _import_class("EmbeddedDocumentField")
+        ReferenceField = _import_class("ReferenceField")
+        ListField = _import_class("ListField")
+
         # only root fields ['test1.a', 'test2'] => ['test1', 'test2']
         root_fields = set([f.split('.')[0] for f in fields])
 
+        if populate:
+            populate_list = [ _field.split('.') for _field in populate.split(',')]
+            populate_domain = dict((_field[0],_field[1:]) for _field in populate_list)
+
         for field_name in self:
+
             if root_fields and field_name not in root_fields:
                 continue
 
@@ -323,22 +333,49 @@ def to_mongo(self, use_db_field=True, fields=None):
 
             if field is None and self._dynamic:
                 field = self._dynamic_fields.get(field_name)
-
+            
             if value is not None:
-                f_inputs = field.to_mongo.__code__.co_varnames
-                ex_vars = {}
-                if fields and 'fields' in f_inputs:
-                    key = '%s.' % field_name
-                    embedded_fields = [
-                        i.replace(key, '') for i in fields
-                        if i.startswith(key)]
+                if populate and field_name in populate_domain.keys():
+                    if isinstance(field, ListField):
+                        _obj = []
+                        for ref in value:
+                            if isinstance(ref, DBRef):
+                                _ref_model = field.field.document_type
+                                
+                                if isinstance(ref.id, ObjectId):
+                                    ref_id = ref.id
+                                    _obj.append(_ref_model.objects.get(id=ref_id).to_mongo(populate='.'.join(populate_domain[field_name])))
+                            elif isinstance(ref, Document):
+                                _obj.append(ref.to_mongo(populate='.'.join(populate_domain[field_name])))
+                            else:
+                                _obj.append(ref)
+                
+                        value = _obj
+
+                    if isinstance(field, ReferenceField) :
+                        _ref_model = field.document_type
+                        try:
+                            _obj = _ref_model.objects.get(id=value.id).to_mongo(populate='.'.join(populate_domain[field_name]))
+                        except:
+                            _obj = None
+
+                        value = _obj
+
+                else :
+                    f_inputs = field.to_mongo.__code__.co_varnames
+                    ex_vars = {}
+                    if fields and 'fields' in f_inputs:
+                        key = '%s.' % field_name
+                        embedded_fields = [
+                            i.replace(key, '') for i in fields
+                            if i.startswith(key)]
 
-                    ex_vars['fields'] = embedded_fields 
+                        ex_vars['fields'] = embedded_fields 
 
-                if 'use_db_field' in f_inputs:
-                    ex_vars['use_db_field'] = use_db_field
-                     
-                value = field.to_mongo(value, **ex_vars)
+                    if 'use_db_field' in f_inputs:
+                        ex_vars['use_db_field'] = use_db_field
+                         
+                    value = field.to_mongo(value, **ex_vars)
 
             # Handle self generating fields
             if value is None and field._auto_gen:
@@ -350,9 +387,9 @@ def to_mongo(self, use_db_field=True, fields=None):
                     data[field.db_field] = value
                 else:
                     data[field.name] = value
+        
 
         # If "_id" has not been set, then try and set it
-        Document = _import_class("Document")
         if isinstance(self, Document):
             if data["_id"] is None:
                 data["_id"] = self._data.get("id", None)
@@ -365,7 +402,8 @@ def to_mongo(self, use_db_field=True, fields=None):
                 not self._meta.get('allow_inheritance', ALLOW_INHERITANCE)):
             data.pop('_cls')
 
-        return data
+        return data    
+
 
     def validate(self, clean=True):
         """Ensure that all fields' values are valid and that required fields
@@ -416,9 +454,11 @@ def to_json(self, *args, **kwargs):
         """Converts a document to JSON.
         :param use_db_field: Set to True by default but enables the output of the json structure with the field names
             and not the mongodb store db_names in case of set to False
+        :param populate: 
         """
+        populate = kwargs.pop('populate', False)
         use_db_field = kwargs.pop('use_db_field', True)
-        return json_util.dumps(self.to_mongo(use_db_field), *args, **kwargs)
+        return json_util.dumps(self.to_mongo(use_db_field, populate=populate), *args, **kwargs)
 
     @classmethod
     def from_json(cls, json_data, created=False):

tests/document/json_serialisation.py

@@ -63,6 +63,28 @@ def __eq__(self, other):
 
         self.assertEqual(doc, Doc.from_json(doc.to_json()))
 
+    def test_populate(self):
+        class VerySimple(Document):
+            num = StringField(default='2')
+
+        class Simple(Document):
+            val = StringField(default='1')
+            inner = ReferenceField(VerySimple)
+
+        class Doc(Document):
+            simple_ref = ReferenceField(Simple)
+            list_ref = ListField(ReferenceField(Simple))
+
+        docum = Doc(simple_ref=Simple(inner=VerySimple().save()).save())
+        docum.list_ref=[Simple(inner=VerySimple().save()).save(), Simple().save()]
+        docum.save()
+        _json = docum.to_json( populate='simple_ref,list_ref.inner')
+        
+        import json
+        _json = json.loads(_json)
+        self.assertEqual(_json['simple_ref']['val'],'1')
+        self.assertEqual(_json['list_ref'][0]['inner']['num'],'2')
+
     def test_json_complex(self):
 
         if pymongo.version_tuple[0] <= 2 and pymongo.version_tuple[1] <= 3: