INTPYTHON-624 Add PolymorphicEmbeddedModelField

Author: timgraham

django_mongodb_backend/compiler.py

@@ -746,7 +746,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,217 @@
+import contextlib
+import difflib
+
+from django.core import checks
+from django.core.exceptions import FieldDoesNotExist, ValidationError
+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 classes once the field knows the model it
+        belongs to. If any of the items in __init__()'s embedded_models
+        argument are strings, resolve each to the actual model class,
+        similar to relation fields.
+        """
+        self._model = model
+        if model is not None:
+            for embedded_model in self.embedded_models:
+                if isinstance(embedded_model, str):
+
+                    def _resolve_lookup(_, *resolved_models):
+                        self.embedded_models = resolved_models
+
+                    lazy_related_operation(_resolve_lookup, model, *self.embedded_models)
+
+    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 = None
+        for model in self.embedded_models:
+            with contextlib.suppress(FieldDoesNotExist):
+                field = model._meta.get_field(name)
+        if field is None:
+            raise FieldDoesNotExist(
+                f"The models of field '{self.name}' have no field named '{name}'."
+            )
+        return KeyTransformFactory(name, field)
+
+    def validate(self, value, model_instance):
+        super().validate(value, model_instance)
+        if not isinstance(value, self.embedded_models):
+            raise ValidationError(
+                f"Expected instance of type {self.embedded_models!r}, not {type(value)!r}."
+            )
+        for field in value._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()

docs/source/ref/models/fields.rst

@@ -313,3 +313,19 @@ These indexes use 0-based indexing.
 .. class:: ObjectIdField
 
     Stores an :class:`~bson.objectid.ObjectId`.
+
+``PolymorphicEmbeddedModelField``
+---------------------------------
+
+.. class:: PolymorphicEmbeddedModelField(embedded_models, **kwargs)
+
+    Stores a model of type ``embedded_models``.
+
+    .. attribute:: embedded_models
+
+        This is a required argument.
+
+        Specifies a list of model classes that may be embedded.
+
+        Each model class reference works just like
+        :attr:`.EmbeddedModelField.embedded_model`.

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 Person(models.Model):
+    name = models.CharField(max_length=100)
+    pet = PolymorphicEmbeddedModelField(("Dog", "Cat"), blank=True, null=True)
+
+    def __str__(self):
+        return self.name
+
+
+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