Improve complexbasefield perf for primtive types

mongoengine/base/document.py

@@ -822,7 +822,9 @@ def _from_son(cls, son, _auto_dereference=True, created=False):
                 value = data[field.db_field]
                 try:
                     data[field_name] = (
-                        value if value is None else field.to_python(value)
+                        value
+                        if value is None
+                        else field._to_python_safe_call(value, bson_native=True)
                     )
                     if field_name != field.db_field:
                         del data[field.db_field]

mongoengine/base/fields.py

@@ -18,6 +18,12 @@
 __all__ = ("BaseField", "ComplexBaseField", "ObjectIdField", "GeoJsonBaseField")
 
 
+# Types whose values do not need any conversion when they appear inside a
+# ComplexBaseField (DictField / ListField). Short-circuiting these avoids
+# re-entering to_python for every primitive leaf of a large nested tree.
+_PRIMITIVE_TYPES = frozenset((str, int, float, bool, type(None), bytes))
+
+
 @contextlib.contextmanager
 def _no_dereference_for_fields(*fields):
     """Context manager for temporarily disabling a Field's auto-dereferencing
@@ -40,6 +46,14 @@ class BaseField:
     _geo_index = False
     _auto_gen = False  # Call `generate` to generate a value
     _thread_local_storage = threading.local()
+    # Whether this field's ``to_python`` accepts the internal ``bson_native``
+    # kwarg used by the MongoDB read path (:meth:`BaseDocument._from_son`).
+    # Kept as a class-level flag so the fast path can dispatch without per-call
+    # signature inspection. Overridden to ``True`` on :class:`ComplexBaseField`
+    # and only subclasses whose ``to_python`` signature declares
+    # ``bson_native`` should set it. Read-only from the field consumer's
+    # perspective.
+    _to_python_accepts_bson_native = False
 
     # These track each time a Field instance is created. Used to retain order.
     # The auto_creation_counter is used for fields that MongoEngine implicitly
@@ -222,6 +236,19 @@ def to_mongo(self, value):
         """Convert a Python type to a MongoDB-compatible type."""
         return self.to_python(value)
 
+    def _to_python_safe_call(self, value, bson_native=False):
+        """Helper method to call to_python, forwarding ``bson_native`` only
+        when the field opts in via :attr:`_to_python_accepts_bson_native`.
+
+        User-defined Field subclasses that override ``to_python(self, value)``
+        without knowing about the internal ``bson_native`` kwarg keep working
+        — passing it blindly would raise ``TypeError``. Mirrors
+        :meth:`_to_mongo_safe_call`.
+        """
+        if self._to_python_accepts_bson_native:
+            return self.to_python(value, bson_native=bson_native)
+        return self.to_python(value)
+
     def _to_mongo_safe_call(self, value, use_db_field=True, fields=None):
         """Helper method to call to_mongo with proper inputs."""
         f_inputs = self.to_mongo.__code__.co_varnames
@@ -313,6 +340,8 @@ class ComplexBaseField(BaseField):
     items in a list / dict rather than one at a time.
     """
 
+    _to_python_accepts_bson_native = True
+
     def __init__(self, field=None, **kwargs):
         if field is not None and not isinstance(field, BaseField):
             raise TypeError(
@@ -403,8 +432,17 @@ def __get__(self, instance, owner):
 
         return value
 
-    def to_python(self, value):
-        """Convert a MongoDB-compatible type to a Python type."""
+    def to_python(self, value, *, bson_native=False):
+        """Convert a MongoDB-compatible type to a Python type.
+
+        ``bson_native=True`` is an internal signal set by :meth:`_from_son`
+        (the MongoDB read path) meaning ``value`` came straight out of
+        pymongo's BSON parser and therefore only contains BSON-native types
+        at every level. In that case there is nothing left to convert on the
+        untyped branch, so the value can be aliased instead of walked. Do not
+        pass ``bson_native=True`` for user-supplied data — it would silently
+        skip Document→DBRef conversion and ``.to_python()`` delegation.
+        """
         if isinstance(value, str):
             return value
 
@@ -426,14 +464,31 @@ def to_python(self, value):
 
         if self.field:
             self.field.set_auto_dereferencing(self._auto_dereference)
-            value_dict = {
-                key: self.field.to_python(item) for key, item in value.items()
-            }
+            field_to_python = self.field.to_python
+            if self.field._to_python_accepts_bson_native:
+                value_dict = {
+                    k: field_to_python(v, bson_native=bson_native)
+                    for k, v in value.items()
+                }
+            else:
+                value_dict = {k: field_to_python(v) for k, v in value.items()}
+        elif bson_native and not is_list:
+            # BSON-native tree with no sub-field: none of the untyped-branch
+            # conversions (Document→DBRef, value.to_python(), nested recursion)
+            # apply to BSON-native values, so aliasing the input is safe and
+            # skips the O(n) rebuild entirely.
+            return value
         else:
             Document = _import_class("Document")
             value_dict = {}
             for k, v in value.items():
-                if isinstance(v, Document):
+                # Primitive leaves (the overwhelming majority of entries in a
+                # dict-of-primitives read from MongoDB) need no conversion.
+                # Short-circuiting them avoids re-entering to_python, which
+                # is where large nested trees spend most of their time.
+                if type(v) in _PRIMITIVE_TYPES:
+                    value_dict[k] = v
+                elif isinstance(v, Document):
                     # We need the id from the saved object to create the DBRef
                     if v.pk is None:
                         self.error(
@@ -445,7 +500,7 @@ def to_python(self, value):
                 elif hasattr(v, "to_python"):
                     value_dict[k] = v.to_python()
                 else:
-                    value_dict[k] = self.to_python(v)
+                    value_dict[k] = self.to_python(v, bson_native=bson_native)
 
         if is_list:  # Convert back to a list
             return [

tests/fields/test_dict_field.py

@@ -1,5 +1,5 @@
 import pytest
-from bson import InvalidDocument
+from bson import DBRef, InvalidDocument, ObjectId
 
 from mongoengine import *
 from mongoengine.base import BaseDict
@@ -10,6 +10,169 @@
 from tests.utils import MongoDBTestCase, get_as_pymongo
 
 
+class TestDictFieldToPython:
+    def test_to_python__untyped_dict_contains_document__converts_to_dbref(self):
+        class Referenced(Document):
+            pass
+
+        referenced = Referenced(id=ObjectId())
+
+        converted = DictField().to_python(
+            {
+                "referenced": referenced,
+                "nested": [{"referenced": referenced}],
+            }
+        )
+
+        assert isinstance(converted["referenced"], DBRef)
+        assert converted["referenced"] == DBRef(
+            Referenced._get_collection_name(), referenced.id
+        )
+        assert isinstance(converted["nested"][0]["referenced"], DBRef)
+
+    def test_to_python__untyped_dict_contains_convertible_value__converts_value(self):
+        class Convertible:
+            def to_python(self):
+                return "converted"
+
+        converted = DictField().to_python({"value": Convertible()})
+
+        assert converted == {"value": "converted"}
+
+    def test_to_python__auto_dereferencing_disabled__propagates_to_nested_field(self):
+        class Referenced(Document):
+            pass
+
+        class Embedded(EmbeddedDocument):
+            referenced = ReferenceField(Referenced)
+
+        field = DictField(EmbeddedDocumentField(Embedded))
+        field.set_auto_dereferencing(False)
+
+        converted = field.to_python({"value": {"referenced": ObjectId()}})
+
+        assert converted["value"]._fields["referenced"]._auto_dereference is False
+
+    def test_to_python__typed_dict_receives_truthy_non_dict__raises_validation_error(
+        self,
+    ):
+        class Model(Document):
+            values = DictField(IntField())
+
+        with pytest.raises(
+            ValidationError, match="Only dictionaries may be used in a DictField"
+        ):
+            Model(values=[1]).validate()
+
+    @pytest.mark.parametrize(
+        "field,value",
+        [
+            pytest.param(DictField(null=True), {}, id="top-level-nullable"),
+            pytest.param(ListField(DictField()), [{}], id="nested-dict-field"),
+            pytest.param(ListField(MapField(IntField())), [{}], id="nested-map-field"),
+        ],
+    )
+    def test_to_python__value_contains_empty_dict__preserves_empty_dict(
+        self, field, value
+    ):
+        assert field.to_python(value) == value
+
+    def test_to_python__untyped_nested_primitives__preserves_shape(self):
+        field = DictField()
+        value = {
+            "numbers": [1, 2],
+            "nested": {"enabled": True, "name": "test"},
+        }
+
+        converted = field.to_python(value)
+
+        assert converted == value
+
+    def test_to_python__untyped_dict_contains_dbref__preserves_dbref(self):
+        oid = ObjectId()
+        dbref = DBRef("collection", oid)
+
+        converted = DictField().to_python({"ref": dbref})
+
+        assert converted == {"ref": dbref}
+        assert isinstance(converted["ref"], DBRef)
+
+    def test_to_python__dict_subclass_with_falsy_bool__preserves_entries(self):
+        """A dict subclass whose __bool__ is False must not be silently dropped."""
+
+        class FalsyDict(dict):
+            def __bool__(self):
+                return False
+
+        value = FalsyDict({"a": "1", "b": "2"})
+        assert not value  # sanity check
+
+        converted = DictField(IntField()).to_python(value)
+
+        assert type(converted) is dict
+        assert converted == {"a": 1, "b": 2}
+
+    def test_to_python__mapfield_typed_dict_of_primitives__preserves_shape(self):
+        """MapField inherits DictField.to_python; primitive dict must round-trip."""
+        converted = MapField(IntField()).to_python({"a": 1, "b": 2})
+
+        assert converted == {"a": 1, "b": 2}
+
+    def test_to_python__mapfield_delegates_to_nested_field(self):
+        class Doubling(IntField):
+            def to_python(self, value):
+                return value * 2
+
+        converted = MapField(Doubling()).to_python({"a": 1, "b": 2})
+
+        assert converted == {"a": 2, "b": 4}
+
+    def test_to_python__mapfield_receives_truthy_non_dict__raises_validation_error(
+        self,
+    ):
+        class Model(Document):
+            values = MapField(IntField())
+
+        with pytest.raises(
+            ValidationError, match="Only dictionaries may be used in a DictField"
+        ):
+            Model(values=[1]).validate()
+
+    def test_to_python__bson_native_untyped_dict__aliases_input(self):
+        """When called with bson_native=True (only from _from_son) the
+        untyped branch aliases the input rather than rebuilding, skipping the
+        O(n) walk.
+        """
+        value = {"a": 1, "b": {"nested": 2}}
+
+        converted = DictField().to_python(value, bson_native=True)
+
+        assert converted is value
+
+    def test_to_python__default_not_bson_native__still_walks_and_converts(self):
+        """The default (no ``bson_native`` kwarg) preserves the pre-existing
+        untyped-branch semantics: Documents become DBRefs, ``.to_python()``
+        gets delegated. User code calling ``field.to_python(value)`` must
+        never accidentally hit the fast path.
+        """
+
+        class Referenced(Document):
+            pass
+
+        class Convertible:
+            def to_python(self):
+                return "converted"
+
+        referenced = Referenced(id=ObjectId())
+        value = {"ref": referenced, "conv": Convertible()}
+
+        converted = DictField().to_python(value)  # no bson_native kwarg
+
+        assert converted is not value
+        assert isinstance(converted["ref"], DBRef)
+        assert converted["conv"] == "converted"
+
+
 class TestDictField(MongoDBTestCase):
     def test_storage(self):
         class BlogPost(Document):
@@ -274,6 +437,24 @@ class MyModel(Document):
         expected_raw_doc = {"_id": doc.id, "events": [{"a": 1}, {}]}
         assert raw_doc == expected_raw_doc
 
+        doc.reload()
+        assert doc.events[-1] == {}
+        assert isinstance(doc.events[-1], dict)
+
+    def test_update__push_empty_dict_to_list_of_mapfield__preserves_dict_on_reload(
+        self,
+    ):
+        class MyModel(Document):
+            events = ListField(MapField(IntField()))
+
+        doc = MyModel(events=[{"a": 1}]).save()
+        MyModel.objects(id=doc.id).update(push__events={})
+
+        doc.reload()
+
+        assert doc.events[-1] == {}
+        assert isinstance(doc.events[-1], dict)
+
     def test_ensure_unique_default_instances(self):
         """Ensure that every field has it's own unique default instance."""
 
@@ -388,3 +569,42 @@ class Simple(Document):
         assert isinstance(s.mapping7["someint"][0]["d"], Doc)
         assert isinstance(s.mapping8["someint"][0]["d"][0], Doc)
         assert isinstance(s.mapping9["someint"][0]["d"][0], Doc)
+
+    def test_dictfield_with_embeddeddocument_field_roundtrip(self):
+        """Ensure DictField(EmbeddedDocumentField) rebuilds the embedded instance."""
+
+        class Setting(EmbeddedDocument):
+            value = StringField()
+
+        class Simple(Document):
+            mapping = DictField(EmbeddedDocumentField(Setting))
+
+        Simple.drop_collection()
+
+        Simple(mapping={"a": Setting(value="foo"), "b": Setting(value="bar")}).save()
+
+        reloaded = Simple.objects.first()
+        assert isinstance(reloaded.mapping["a"], Setting)
+        assert isinstance(reloaded.mapping["b"], Setting)
+        assert reloaded.mapping["a"].value == "foo"
+        assert reloaded.mapping["b"].value == "bar"
+
+    def test_dictfield_reads_non_dict_stored_in_db_schema_drift(self):
+        """Reading a document whose DB value is not a dict must not blow up.
+
+        Data written by a different tool or an older schema may end up with a
+        non-dict value on a DictField. Preserving the current tolerant behavior
+        avoids breaking existing systems when the field's ``to_python`` is
+        optimized.
+        """
+
+        class Model(Document):
+            m = DictField(field=IntField())
+
+        Model.drop_collection()
+
+        Model._get_collection().insert_one({"_id": 1, "m": [{"a": 1}]})
+        Model._get_collection().insert_one({"_id": 2, "m": "some-string"})
+
+        loaded = {doc.id: doc.m for doc in Model.objects.order_by("id")}
+        assert loaded == {1: [{"a": 1}], 2: "some-string"}