Merge pull request #2488 from bagerard/add_lazy_loading_method_ease_profiling

Add _lazy_load_ref methods to make profiling easier

mongoengine/base/fields.py

@@ -267,6 +267,17 @@ class ComplexBaseField(BaseField):
         self.field = field
         super().__init__(**kwargs)
 
+    @staticmethod
+    def _lazy_load_refs(instance, name, ref_values, *, max_depth):
+        _dereference = _import_class("DeReference")()
+        documents = _dereference(
+            ref_values,
+            max_depth=max_depth,
+            instance=instance,
+            name=name,
+        )
+        return documents
+
     def __get__(self, instance, owner):
         """Descriptor to automatically dereference references."""
         if instance is None:
@@ -284,19 +295,15 @@ class ComplexBaseField(BaseField):
             or isinstance(self.field, (GenericReferenceField, ReferenceField))
         )
 
-        _dereference = _import_class("DeReference")()
-
         if (
             instance._initialised
             and dereference
             and instance._data.get(self.name)
             and not getattr(instance._data[self.name], "_dereferenced", False)
         ):
-            instance._data[self.name] = _dereference(
+            ref_values = instance._data.get(self.name)
-                instance._data.get(self.name),
+            instance._data[self.name] = self._lazy_load_refs(
-                max_depth=1,
+                ref_values=ref_values, instance=instance, name=self.name, max_depth=1
-                instance=instance,
-                name=self.name,
             )
             if hasattr(instance._data[self.name], "_dereferenced"):
                 instance._data[self.name]._dereferenced = True
@@ -322,7 +329,9 @@ class ComplexBaseField(BaseField):
             and isinstance(value, (BaseList, BaseDict))
             and not value._dereferenced
         ):
-            value = _dereference(value, max_depth=1, instance=instance, name=self.name)
+            value = self._lazy_load_refs(
+                ref_values=value, instance=instance, name=self.name, max_depth=1
+            )
             value._dereferenced = True
             instance._data[self.name] = value
 

mongoengine/fields.py

@@ -1194,6 +1194,14 @@ class ReferenceField(BaseField):
                 self.document_type_obj = get_document(self.document_type_obj)
         return self.document_type_obj
 
+    @staticmethod
+    def _lazy_load_ref(ref_cls, dbref):
+        dereferenced_son = ref_cls._get_db().dereference(dbref)
+        if dereferenced_son is None:
+            raise DoesNotExist(f"Trying to dereference unknown document {dbref}")
+
+        return ref_cls._from_son(dereferenced_son)
+
     def __get__(self, instance, owner):
         """Descriptor to allow lazy dereferencing."""
         if instance is None:
@@ -1201,20 +1209,17 @@ class ReferenceField(BaseField):
             return self
 
         # Get value from document instance if available
-        value = instance._data.get(self.name)
+        ref_value = instance._data.get(self.name)
         auto_dereference = instance._fields[self.name]._auto_dereference
         # Dereference DBRefs
-        if auto_dereference and isinstance(value, DBRef):
+        if auto_dereference and isinstance(ref_value, DBRef):
-            if hasattr(value, "cls"):
+            if hasattr(ref_value, "cls"):
                 # Dereference using the class type specified in the reference
-                cls = get_document(value.cls)
+                cls = get_document(ref_value.cls)
             else:
                 cls = self.document_type
-            dereferenced = cls._get_db().dereference(value)
+
-            if dereferenced is None:
+            instance._data[self.name] = self._lazy_load_ref(cls, ref_value)
-                raise DoesNotExist("Trying to dereference unknown document %s" % value)
-            else:
-                instance._data[self.name] = cls._from_son(dereferenced)
 
         return super().__get__(instance, owner)
 
@@ -1353,6 +1358,14 @@ class CachedReferenceField(BaseField):
                 self.document_type_obj = get_document(self.document_type_obj)
         return self.document_type_obj
 
+    @staticmethod
+    def _lazy_load_ref(ref_cls, dbref):
+        dereferenced_son = ref_cls._get_db().dereference(dbref)
+        if dereferenced_son is None:
+            raise DoesNotExist(f"Trying to dereference unknown document {dbref}")
+
+        return ref_cls._from_son(dereferenced_son)
+
     def __get__(self, instance, owner):
         if instance is None:
             # Document class being used rather than a document object
@@ -1364,11 +1377,7 @@ class CachedReferenceField(BaseField):
 
         # Dereference DBRefs
         if auto_dereference and isinstance(value, DBRef):
-            dereferenced = self.document_type._get_db().dereference(value)
+            instance._data[self.name] = self._lazy_load_ref(self.document_type, value)
-            if dereferenced is None:
-                raise DoesNotExist("Trying to dereference unknown document %s" % value)
-            else:
-                instance._data[self.name] = self.document_type._from_son(dereferenced)
 
         return super().__get__(instance, owner)
 
@@ -1493,6 +1502,14 @@ class GenericReferenceField(BaseField):
             value = value._class_name
         super()._validate_choices(value)
 
+    @staticmethod
+    def _lazy_load_ref(ref_cls, dbref):
+        dereferenced_son = ref_cls._get_db().dereference(dbref)
+        if dereferenced_son is None:
+            raise DoesNotExist(f"Trying to dereference unknown document {dbref}")
+
+        return ref_cls._from_son(dereferenced_son)
+
     def __get__(self, instance, owner):
         if instance is None:
             return self
@@ -1500,12 +1517,9 @@ class GenericReferenceField(BaseField):
         value = instance._data.get(self.name)
 
         auto_dereference = instance._fields[self.name]._auto_dereference
-        if auto_dereference and isinstance(value, (dict, SON)):
+        if auto_dereference and isinstance(value, dict):
-            dereferenced = self.dereference(value)
+            doc_cls = get_document(value["_cls"])
-            if dereferenced is None:
+            instance._data[self.name] = self._lazy_load_ref(doc_cls, value["_ref"])
-                raise DoesNotExist("Trying to dereference unknown document %s" % value)
-            else:
-                instance._data[self.name] = dereferenced
 
         return super().__get__(instance, owner)
 
@@ -1524,14 +1538,6 @@ class GenericReferenceField(BaseField):
                 " saved to the database"
             )
 
-    def dereference(self, value):
-        doc_cls = get_document(value["_cls"])
-        reference = value["_ref"]
-        doc = doc_cls._get_db().dereference(reference)
-        if doc is not None:
-            doc = doc_cls._from_son(doc)
-        return doc
-
     def to_mongo(self, document):
         if document is None:
             return None