INTPYTHON-624 Add PolymorphicEmbeddedModelField

Author: timgraham

django_mongodb_backend/compiler.py

@@ -743,7 +743,7 @@ def execute_sql(self, result_type):
             elif hasattr(value, "prepare_database_save"):
                 if field.remote_field:
                     value = value.prepare_database_save(field)
-                elif not hasattr(field, "embedded_model"):
+                elif not (hasattr(field, "embedded_model") or hasattr(field, "embedded_models")):
                     raise TypeError(
                         f"Tried to update field {field} with a model "
                         f"instance, {value!r}. Use a value compatible with "

django_mongodb_backend/fields/__init__.py

@@ -5,6 +5,7 @@
 from .embedded_model_array import EmbeddedModelArrayField
 from .json import register_json_field
 from .objectid import ObjectIdField
+from .polymorphic_embedded_model import PolymorphicEmbeddedModelField
 
 __all__ = [
     "register_fields",
@@ -13,6 +14,7 @@
     "EmbeddedModelField",
     "ObjectIdAutoField",
     "ObjectIdField",
+    "PolymorphicEmbeddedModelField",
 ]
 
 

django_mongodb_backend/fields/polymorphic_embedded_model.py

@@ -0,0 +1,202 @@
+from django.core import checks
+from django.db import models
+
+# from django.db.models.fields.related import lazy_related_operation
+# from django.db.models.lookups import Transform
+
+
+class PolymorphicEmbeddedModelField(models.Field):
+    """Field that stores a model instance."""
+
+    def __init__(self, embedded_models, *args, **kwargs):
+        """
+        `embedded_models` is a list of possible model classes to be stored.
+        Like other relational fields, each model may also be passed as a
+        string.
+        """
+        self.embedded_models = embedded_models
+        super().__init__(*args, **kwargs)
+
+    def db_type(self, connection):
+        return "embeddedDocuments"
+
+    def check(self, **kwargs):
+        from ..models import EmbeddedModel
+
+        errors = super().check(**kwargs)
+        for model in self.embedded_models:
+            if not issubclass(model, EmbeddedModel):
+                return [
+                    checks.Error(
+                        "Embedded models must be a subclass of "
+                        "django_mongodb_backend.models.EmbeddedModel.",
+                        obj=self,
+                        hint="{model} doesn't subclass EmbeddedModel.",
+                        id="django_mongodb_backend.embedded_model.E002",
+                    )
+                ]
+            for field in model._meta.fields:
+                if field.remote_field:
+                    errors.append(
+                        checks.Error(
+                            "Embedded models cannot have relational fields "
+                            f"({model().__class__.__name__}.{field.name} "
+                            f"is a {field.__class__.__name__}).",
+                            obj=self,
+                            id="django_mongodb_backend.embedded_model.E001",
+                        )
+                    )
+        return errors
+
+    def deconstruct(self):
+        name, path, args, kwargs = super().deconstruct()
+        if path.startswith("django_mongodb_backend.fields.polymorphic_embedded_model"):
+            path = path.replace(
+                "django_mongodb_backend.fields.polymorphic_embedded_model",
+                "django_mongodb_backend.fields",
+            )
+        kwargs["embedded_models"] = self.embedded_models
+        return name, path, args, kwargs
+
+    def get_internal_type(self):
+        return "PolymorphicEmbeddedModelField"
+
+    #    def _set_model(self, model):
+    #        """
+    #        Resolve embedded model class once the field knows the model it belongs
+    #        to. If __init__()'s embedded_model argument is a string, resolve it to
+    #        the actual model class, similar to relation fields.
+    #        """
+    #        self._model = model
+    #        if model is not None and isinstance(self.embedded_model, str):
+
+    #            def _resolve_lookup(_, resolved_model):
+    #                self.embedded_model = resolved_model
+
+    #            lazy_related_operation(_resolve_lookup, model, self.embedded_model)
+
+    #    model = property(lambda self: self._model, _set_model)
+
+    def from_db_value(self, value, expression, connection):
+        return self.to_python(value)
+
+    def to_python(self, value):
+        """
+        Pass embedded model fields' values through each field's to_python() and
+        reinstantiate the embedded instance.
+        """
+        if value is None:
+            return None
+        if not isinstance(value, dict):
+            return value
+        model_class = self._get_model_from_label(value.pop("_label"))
+        instance = model_class(
+            **{
+                field.attname: field.to_python(value[field.attname])
+                for field in model_class._meta.fields
+                if field.attname in value
+            }
+        )
+        instance._state.adding = False
+        return instance
+
+    def get_db_prep_save(self, embedded_instance, connection):
+        """
+        Apply pre_save() and get_db_prep_save() of embedded instance fields and
+        create the {field: value} dict to be saved.
+        """
+        if embedded_instance is None:
+            return None
+        if not isinstance(embedded_instance, self.embedded_models):
+            raise TypeError(
+                f"Expected instance of type {self.embedded_models!r}, not "
+                f"{type(embedded_instance)!r}."
+            )
+        field_values = {}
+        add = embedded_instance._state.adding
+        for field in embedded_instance._meta.fields:
+            value = field.get_db_prep_save(
+                field.pre_save(embedded_instance, add), connection=connection
+            )
+            # Exclude unset primary keys (e.g. {'id': None}).
+            if field.primary_key and value is None:
+                continue
+            field_values[field.attname] = value
+        field_values["_label"] = embedded_instance._meta.label
+        # This instance will exist in the database soon.
+        embedded_instance._state.adding = False
+        return field_values
+
+    #    def get_transform(self, name):
+    #        transform = super().get_transform(name)
+    #        if transform:
+    #            return transform
+    #        field = self.embedded_model._meta.get_field(name)
+    #        return KeyTransformFactory(name, field)
+
+    #    def validate(self, value, model_instance):
+    #        super().validate(value, model_instance)
+    #        if self.embedded_model is None:
+    #            return
+    #        for field in self.embedded_model._meta.fields:
+    #            attname = field.attname
+    #            field.validate(getattr(value, attname), model_instance)
+
+    def formfield(self, **kwargs):
+        raise NotImplementedError("PolymorphicEmbeddedModelField does not support forms.")
+
+    def _get_model_from_label(self, label):
+        return {model._meta.label: model for model in self.embedded_models}[label]
+
+
+# class KeyTransform(Transform):
+#    def __init__(self, key_name, ref_field, *args, **kwargs):
+#        super().__init__(*args, **kwargs)
+#        self.key_name = str(key_name)
+#        self.ref_field = ref_field
+
+#    def get_lookup(self, name):
+#        return self.ref_field.get_lookup(name)
+
+#    def get_transform(self, name):
+#        """
+#        Validate that `name` is either a field of an embedded model or a
+#        lookup on an embedded model's field.
+#        """
+#        if transform := self.ref_field.get_transform(name):
+#            return transform
+#        suggested_lookups = difflib.get_close_matches(name, self.ref_field.get_lookups())
+#        if suggested_lookups:
+#            suggested_lookups = " or ".join(suggested_lookups)
+#            suggestion = f", perhaps you meant {suggested_lookups}?"
+#        else:
+#            suggestion = "."
+#        raise FieldDoesNotExist(
+#            f"Unsupported lookup '{name}' for "
+#            f"{self.ref_field.__class__.__name__} '{self.ref_field.name}'"
+#            f"{suggestion}"
+#        )
+
+#    def as_mql(self, compiler, connection):
+#        previous = self
+#        key_transforms = []
+#        while isinstance(previous, KeyTransform):
+#            key_transforms.insert(0, previous.key_name)
+#            previous = previous.lhs
+#        mql = previous.as_mql(compiler, connection)
+#        for key in key_transforms:
+#            mql = {"$getField": {"input": mql, "field": key}}
+#        return mql
+
+#    @property
+#    def output_field(self):
+#        return self.ref_field
+
+
+# class KeyTransformFactory:
+#    def __init__(self, key_name, ref_field):
+#        self.key_name = key_name
+#        self.ref_field = ref_field
+
+#    def __call__(self, *args, **kwargs):
+#        return KeyTransform(self.key_name, self.ref_field, *args, **kwargs)

django_mongodb_backend/operations.py

@@ -122,6 +122,8 @@ def get_db_converters(self, expression):
             )
         elif internal_type == "JSONField":
             converters.append(self.convert_jsonfield_value)
+        elif internal_type == "PolymorphicEmbeddedModelField":
+            converters.append(self.convert_polymorphicembeddedmodelfield_value)
         elif internal_type == "TimeField":
             # Trunc(... output_field="TimeField") values must remain datetime
             # until Trunc.convert_value() so they can be converted from UTC
@@ -182,6 +184,19 @@ def convert_jsonfield_value(self, value, expression, connection):
         """
         return json.dumps(value)
 
+    def convert_polymorphicembeddedmodelfield_value(self, value, expression, connection):
+        if value is not None:
+            model_class = expression.output_field._get_model_from_label(value["_label"])
+            # Apply database converters to each field of the embedded model.
+            for field in model_class._meta.fields:
+                field_expr = Expression(output_field=field)
+                converters = connection.ops.get_db_converters(
+                    field_expr
+                ) + field_expr.get_db_converters(connection)
+                for converter in converters:
+                    value[field.attname] = converter(value[field.attname], field_expr, connection)
+        return value
+
     def convert_timefield_value(self, value, expression, connection):
         if value is not None:
             value = value.time()

tests/model_fields_/models.py

@@ -7,6 +7,7 @@
     EmbeddedModelArrayField,
     EmbeddedModelField,
     ObjectIdField,
+    PolymorphicEmbeddedModelField,
 )
 from django_mongodb_backend.models import EmbeddedModel
 
@@ -222,3 +223,30 @@ class Tour(models.Model):
 
     def __str__(self):
         return f"Tour by {self.guide}"
+
+
+# PolymorphicEmbeddedModelField
+class Dog(EmbeddedModel):
+    name = models.CharField(max_length=100)
+    barks = models.BooleanField(default=True)
+    data = models.JSONField(default=dict)
+
+    def __str__(self):
+        return self.name
+
+
+class Cat(EmbeddedModel):
+    name = models.CharField(max_length=100)
+    purs = models.BooleanField(default=True)
+    weight = models.DecimalField(max_digits=4, decimal_places=2, blank=True, null=True)
+
+    def __str__(self):
+        return self.name
+
+
+class Person(models.Model):
+    name = models.CharField(max_length=100)
+    pet = PolymorphicEmbeddedModelField((Dog, Cat), blank=True, null=True)
+
+    def __str__(self):
+        return self.name