Improve LazyReferenceField and GenericLazyReferenceField with nested fields

mongoengine/fields.py

@@ -28,6 +28,7 @@ except ImportError:
 from mongoengine.base import (BaseDocument, BaseField, ComplexBaseField,
                               GeoJsonBaseField, LazyReference, ObjectIdField,
                               get_document)
+from mongoengine.common import _import_class
 from mongoengine.connection import DEFAULT_CONNECTION_NAME, get_db
 from mongoengine.document import Document, EmbeddedDocument
 from mongoengine.errors import DoesNotExist, InvalidQueryError, ValidationError
@@ -789,6 +790,17 @@ class ListField(ComplexBaseField):
         kwargs.setdefault('default', lambda: [])
         super(ListField, self).__init__(**kwargs)
 
+    def __get__(self, instance, owner):
+        if instance is None:
+            # Document class being used rather than a document object
+            return self
+        value = instance._data.get(self.name)
+        LazyReferenceField = _import_class('LazyReferenceField')
+        GenericLazyReferenceField = _import_class('GenericLazyReferenceField')
+        if isinstance(self.field, (LazyReferenceField, GenericLazyReferenceField)) and value:
+            instance._data[self.name] = [self.field.build_lazyref(x) for x in value]
+        return super(ListField, self).__get__(instance, owner)
+
     def validate(self, value):
         """Make sure that a list of valid fields is being used."""
         if (not isinstance(value, (list, tuple, QuerySet)) or
@@ -2211,17 +2223,10 @@ class LazyReferenceField(BaseField):
                 self.document_type_obj = get_document(self.document_type_obj)
         return self.document_type_obj
 
-    def __get__(self, instance, owner):
-        """Descriptor to allow lazy dereferencing."""
-        if instance is None:
-            # Document class being used rather than a document object
-            return self
-
-        value = instance._data.get(self.name)
+    def build_lazyref(self, value):
         if isinstance(value, LazyReference):
             if value.passthrough != self.passthrough:
-                instance._data[self.name] = LazyReference(
-                    value.document_type, value.pk, passthrough=self.passthrough)
+                value = LazyReference(value.document_type, value.pk, passthrough=self.passthrough)
         elif value is not None:
             if isinstance(value, self.document_type):
                 value = LazyReference(self.document_type, value.pk, passthrough=self.passthrough)
@@ -2230,6 +2235,16 @@ class LazyReferenceField(BaseField):
             else:
                 # value is the primary key of the referenced document
                 value = LazyReference(self.document_type, value, passthrough=self.passthrough)
+        return value
+
+    def __get__(self, instance, owner):
+        """Descriptor to allow lazy dereferencing."""
+        if instance is None:
+            # Document class being used rather than a document object
+            return self
+
+        value = self.build_lazyref(instance._data.get(self.name))
+        if value:
             instance._data[self.name] = value
 
         return super(LazyReferenceField, self).__get__(instance, owner)
@@ -2254,7 +2269,7 @@ class LazyReferenceField(BaseField):
 
     def validate(self, value):
         if isinstance(value, LazyReference):
-            if not issubclass(value.document_type, self.document_type):
+            if value.collection != self.document_type._get_collection_name():
                 self.error('Reference must be on a `%s` document.' % self.document_type)
             pk = value.pk
         elif isinstance(value, self.document_type):
@@ -2314,23 +2329,26 @@ class GenericLazyReferenceField(GenericReferenceField):
 
     def _validate_choices(self, value):
         if isinstance(value, LazyReference):
-            value = value.document_type
+            value = value.document_type._class_name
         super(GenericLazyReferenceField, self)._validate_choices(value)
 
-    def __get__(self, instance, owner):
-        if instance is None:
-            return self
-
-        value = instance._data.get(self.name)
+    def build_lazyref(self, value):
         if isinstance(value, LazyReference):
             if value.passthrough != self.passthrough:
-                instance._data[self.name] = LazyReference(
-                    value.document_type, value.pk, passthrough=self.passthrough)
+                value = LazyReference(value.document_type, value.pk, passthrough=self.passthrough)
         elif value is not None:
             if isinstance(value, (dict, SON)):
                 value = LazyReference(get_document(value['_cls']), value['_ref'].id, passthrough=self.passthrough)
             elif isinstance(value, Document):
                 value = LazyReference(type(value), value.pk, passthrough=self.passthrough)
+        return value
+
+    def __get__(self, instance, owner):
+        if instance is None:
+            return self
+
+        value = self.build_lazyref(instance._data.get(self.name))
+        if value:
             instance._data[self.name] = value
 
         return super(GenericLazyReferenceField, self).__get__(instance, owner)
@@ -2348,7 +2366,7 @@ class GenericLazyReferenceField(GenericReferenceField):
         if isinstance(document, LazyReference):
             return SON((
                 ('_cls', document.document_type._class_name),
-                ('_ref', document)
+                ('_ref', DBRef(document.document_type._get_collection_name(), document.pk))
             ))
         else:
             return super(GenericLazyReferenceField, self).to_mongo(document)