INTPYTHON-658 Add PolymorphicEmbeddedModelArrayField

Author: timgraham

django_mongodb_backend/fields/__init__.py

@@ -6,6 +6,7 @@
 from .json import register_json_field
 from .objectid import ObjectIdField
 from .polymorphic_embedded_model import PolymorphicEmbeddedModelField
+from .polymorphic_embedded_model_array import PolymorphicEmbeddedModelArrayField
 
 __all__ = [
     "register_fields",
@@ -15,6 +16,7 @@
     "ObjectIdAutoField",
     "ObjectIdField",
     "PolymorphicEmbeddedModelField",
+    "PolymorphicEmbeddedModelArrayField",
 ]
 
 

django_mongodb_backend/fields/embedded_model_array.py

@@ -224,6 +224,8 @@ class EmbeddedModelArrayFieldLessThanOrEqual(
 
 
 class KeyTransform(Transform):
+    field_class_name = "EmbeddedModelArrayField"
+
     def __init__(self, key_name, array_field, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.array_field = array_field
@@ -269,7 +271,7 @@ def get_transform(self, name):
             suggestion = ""
         raise FieldDoesNotExist(
             f"Unsupported lookup '{name}' for "
-            f"EmbeddedModelArrayField of '{output_field.__class__.__name__}'"
+            f"{self.field_class_name} of '{output_field.__class__.__name__}'"
             f"{suggestion}"
         )
 

django_mongodb_backend/fields/polymorphic_embedded_model_array.py

@@ -0,0 +1,116 @@
+import contextlib
+
+from django.core.exceptions import FieldDoesNotExist
+from django.db.models.expressions import Col
+from django.db.models.fields.related import lazy_related_operation
+from django.db.models.lookups import Lookup, Transform
+
+from . import PolymorphicEmbeddedModelField
+from .array import ArrayField, ArrayLenTransform
+from .embedded_model_array import KeyTransform as ArrayFieldKeyTransform
+from .embedded_model_array import KeyTransformFactory as ArrayFieldKeyTransformFactory
+
+
+class PolymorphicEmbeddedModelArrayField(ArrayField):
+    value_is_model_instance = True
+
+    def __init__(self, embedded_models, **kwargs):
+        if "size" in kwargs:
+            raise ValueError("PolymorphicEmbeddedModelArrayField does not support size.")
+        kwargs["editable"] = False
+        super().__init__(PolymorphicEmbeddedModelField(embedded_models), **kwargs)
+        self.embedded_models = embedded_models
+
+    def contribute_to_class(self, cls, name, private_only=False, **kwargs):
+        super().contribute_to_class(cls, name, private_only=private_only, **kwargs)
+
+        if not cls._meta.abstract:
+            # If the embedded_model argument is a string, resolve it to the
+            # actual model class.
+            def _resolve_lookup(_, *resolved_models):
+                self.embedded_models = resolved_models
+
+            lazy_related_operation(_resolve_lookup, cls, *self.embedded_models)
+
+    def deconstruct(self):
+        name, path, args, kwargs = super().deconstruct()
+        if path == (
+            "django_mongodb_backend.fields.polymorphic_embedded_model_array."
+            "PolymorphicEmbeddedModelArrayField"
+        ):
+            path = "django_mongodb_backend.fields.PolymorphicEmbeddedModelArrayField"
+        kwargs["embedded_models"] = self.embedded_models
+        del kwargs["base_field"]
+        del kwargs["editable"]
+        return name, path, args, kwargs
+
+    def get_db_prep_value(self, value, connection, prepared=False):
+        if isinstance(value, list | tuple):
+            # Must call get_db_prep_save() rather than get_db_prep_value()
+            # to transform model instances to dicts.
+            return [self.base_field.get_db_prep_save(i, connection) for i in value]
+        if value is not None:
+            raise TypeError(
+                f"Expected list of {self.embedded_models!r} instances, not {type(value)!r}."
+            )
+        return value
+
+    def formfield(self, **kwargs):
+        raise NotImplementedError("PolymorphicEmbeddedModelField does not support forms.")
+
+    _get_model_from_label = PolymorphicEmbeddedModelField._get_model_from_label
+
+    def get_transform(self, name):
+        transform = super().get_transform(name)
+        if transform:
+            return transform
+        return KeyTransformFactory(name, self)
+
+    def _get_lookup(self, lookup_name):
+        lookup = super()._get_lookup(lookup_name)
+        if lookup is None or lookup is ArrayLenTransform:
+            return lookup
+
+        class EmbeddedModelArrayFieldLookups(Lookup):
+            def as_mql(self, compiler, connection):
+                raise ValueError(
+                    "Lookups aren't supported on PolymorphicEmbeddedModelArrayField. "
+                    "Try querying one of its embedded fields instead."
+                )
+
+        return EmbeddedModelArrayFieldLookups
+
+
+class KeyTransform(ArrayFieldKeyTransform):
+    field_class_name = "PolymorphicEmbeddedModelArrayField"
+
+    def __init__(self, key_name, array_field, *args, **kwargs):
+        # Skip ArrayFieldKeyTransform.__init__()
+        Transform.__init__(self, *args, **kwargs)
+        self.array_field = array_field
+        self.key_name = key_name
+        # Lookups iterate over the array of embedded models. A virtual column
+        # of the queried field's type represents each element.
+        for model in array_field.base_field.embedded_models:
+            with contextlib.suppress(FieldDoesNotExist):
+                field = model._meta.get_field(key_name)
+                break
+        else:
+            raise FieldDoesNotExist(
+                f"The models of field '{array_field.name}' have no field named '{key_name}'."
+            )
+        column_target = field.clone()
+        column_name = f"$item.{key_name}"
+        column_target.db_column = column_name
+        column_target.set_attributes_from_name(column_name)
+        self._lhs = Col(None, column_target)
+        self._sub_transform = None
+
+
+class KeyTransformFactory(ArrayFieldKeyTransformFactory):
+    def __init__(self, key_name, base_field):
+        self.key_name = key_name
+        self.base_field = base_field
+
+    def __call__(self, *args, **kwargs):
+        return KeyTransform(self.key_name, self.base_field, *args, **kwargs)

django_mongodb_backend/operations.py

@@ -124,6 +124,15 @@ def get_db_converters(self, expression):
             converters.append(self.convert_jsonfield_value)
         elif internal_type == "PolymorphicEmbeddedModelField":
             converters.append(self.convert_polymorphicembeddedmodelfield_value)
+        elif internal_type == "PolymorphicEmbeddedModelArrayField":
+            converters.extend(
+                [
+                    self._get_arrayfield_converter(converter)
+                    for converter in self.get_db_converters(
+                        Expression(output_field=expression.output_field.base_field)
+                    )
+                ]
+            )
         elif internal_type == "TimeField":
             # Trunc(... output_field="TimeField") values must remain datetime
             # until Trunc.convert_value() so they can be converted from UTC

docs/source/ref/models/fields.rst

@@ -343,3 +343,39 @@ These indexes use 0-based indexing.
 .. admonition:: Forms are not supported
 
     ``PolymorphicEmbeddedModelField``\s don't appear in model forms.
+
+``PolymorphicEmbeddedModelArrayField``
+--------------------------------------
+
+.. class:: PolymorphicEmbeddedModelArrayField(embedded_models, **kwargs)
+
+    .. versionadded:: 5.2.0b2
+
+    Similar to :class:`PolymorphicEmbeddedModelField`, but stores a **list** of
+    models of type ``embedded_models`` rather than a single instance.
+
+    .. attribute:: embedded_models
+
+        This is a required argument that works just like
+        :attr:`PolymorphicEmbeddedModelField.embedded_models`.
+
+    .. attribute:: max_size
+
+        This is an optional argument.
+
+        If passed, the list will have a maximum size as specified, validated
+        by forms and model validation, but not enforced by the database.
+
+    See :ref:`the embedded model topic guide
+    <polymorphic-embedded-model-array-field-example>` for more details and
+    examples.
+
+.. admonition:: Migrations support is limited
+
+    :djadmin:`makemigrations` does not yet detect changes to embedded models,
+    nor does it create indexes or constraints for embedded models referenced
+    by ``PolymorphicEmbeddedModelArrayField``.
+
+.. admonition:: Forms are not supported
+
+    ``PolymorphicEmbeddedModelArrayField``\s don't appear in model forms.

docs/source/releases/5.2.x.rst

@@ -14,8 +14,9 @@ New features
 - Added the ``options`` parameter to
   :func:`~django_mongodb_backend.utils.parse_uri`.
 - Added support for :ref:`database transactions <transactions>`.
-- Added :class:`~.fields.PolymorphicEmbeddedModelField` for storing a model
-  instance that may be of more than one model class.
+- Added :class:`~.fields.PolymorphicEmbeddedModelField` and
+  :class:`~.fields.PolymorphicEmbeddedModelArrayField` for storing a model
+  instance or list of model instances that may be of more than one model class.
 
 5.2.0 beta 1
 ============

tests/model_fields_/models.py

@@ -7,6 +7,7 @@
     EmbeddedModelArrayField,
     EmbeddedModelField,
     ObjectIdField,
+    PolymorphicEmbeddedModelArrayField,
     PolymorphicEmbeddedModelField,
 )
 from django_mongodb_backend.models import EmbeddedModel
@@ -240,6 +241,7 @@ class Dog(EmbeddedModel):
     created_at = models.DateTimeField(auto_now_add=True)
     updated_at = models.DateTimeField(auto_now=True)
     favorite_toy = PolymorphicEmbeddedModelField(["Bone"], blank=True, null=True)
+    toys = PolymorphicEmbeddedModelArrayField(["Bone"], blank=True, null=True)
 
     def __str__(self):
         return self.name
@@ -250,6 +252,7 @@ class Cat(EmbeddedModel):
     purs = models.BooleanField(default=True)
     weight = models.DecimalField(max_digits=4, decimal_places=2, blank=True, null=True)
     favorite_toy = PolymorphicEmbeddedModelField(["Mouse"], blank=True, null=True)
+    toys = PolymorphicEmbeddedModelArrayField(["Mouse"], blank=True, null=True)
 
     def __str__(self):
         return self.name
@@ -267,3 +270,12 @@ class Mouse(EmbeddedModel):
 
     def __str__(self):
         return self.manufacturer
+
+
+# PolymorphicEmbeddedModelArrayField
+class Owner(models.Model):
+    name = models.CharField(max_length=100)
+    pets = PolymorphicEmbeddedModelArrayField(("Dog", "Cat"), blank=True, null=True)
+
+    def __str__(self):
+        return self.name