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,110 @@
+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):
+    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
+        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}'."
+            )
+        # Lookups iterate over the array of embedded models. A virtual column
+        # of the queried field's type represents each element.
+        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 __call__(self, *args, **kwargs):
+        return KeyTransform(self.key_name, self.base_field, *args, **kwargs)

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
 ============

docs/source/topics/embedded-models.rst

@@ -116,6 +116,8 @@ Represented in BSON, the post's structure looks like this:
       tags: [ { name: 'welcome' }, { name: 'test' } ]
     }
 
+.. _querying-embedded-model-array-field:
+
 Querying ``EmbeddedModelArrayField``
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
@@ -264,6 +266,8 @@ bark::
 
     >>> Person.objects.filter(pet__barks=True)
 
+.. _polymorphic-embedded-model-field-clashing-field-names:
+
 Clashing field names
 ~~~~~~~~~~~~~~~~~~~~
 
@@ -292,3 +296,87 @@ field it finds, ``Target1.number`` in this case.
 Similarly, querying into nested embedded model fields with the same name isn't
 well supported: the first model in ``embedded_models`` is the one that will be
 used for nested lookups.
+
+.. _polymorphic-embedded-model-array-field-example:
+
+``PolymorphicEmbeddedModelArrayField``
+--------------------------------------
+
+The basics
+~~~~~~~~~~
+
+Let's consider this example::
+
+    from django.db import models
+
+    from django_mongodb_backend.fields import PolymorphicEmbeddedModelArrayField
+    from django_mongodb_backend.models import EmbeddedModel
+
+
+    class Person(models.Model):
+        name = models.CharField(max_length=255)
+        pets = PolymorphicEmbeddedModelArrayField(["Cat", "Dog"])
+
+        def __str__(self):
+            return self.name
+
+
+    class Cat(EmbeddedModel):
+        name = models.CharField(max_length=255)
+        purrs = models.BooleanField(default=True)
+
+        def __str__(self):
+            return self.name
+
+
+    class Dog(EmbeddedModel):
+        name = models.CharField(max_length=255)
+        barks = models.BooleanField(default=True)
+
+        def __str__(self):
+            return self.name
+
+
+The API is similar to that of Django's relational fields::
+
+    >>> bob = Person.objects.create(
+    ...     name="Bob",
+    ...     pets=[Dog(name="Woofer"), Cat(name="Phoebe")],
+    ... )
+    >>> bob.pets
+    [<Dog: Woofer>, <Cat: Phoebe>]
+    >>> bob.pets[0].name
+    'Woofer'
+
+Represented in BSON, the Bob's structure looks like this:
+
+.. code-block:: js
+
+    {
+      _id: ObjectId('6875605cf6dc6f95cadf2d75'),
+      name: 'Bob',
+      pets: [
+        { name: 'Woofer', barks: true, _label: 'polymorphic_array.Dog' },
+        { name: 'Phoebe', purrs: true, _label: 'polymorphic_array.Cat' }
+      ]
+    }
+
+The ``_label`` field tracks each model's :attr:`~django.db.models.Options.label`
+so that the models can be initialized properly.
+
+Querying ``PolymorphicEmbeddedModelArrayField``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+You can query into an embedded model array using :ref:`the same syntax and operators
+<querying-embedded-model-array-field>` as :class:`~.fields.EmbeddedModelArrayField`.
+
+Like :class:`~.fields.PolymorphicEmbeddedModelField`, if you filter on fields that aren't shared
+among the embedded models, you'll only get back objects that have embedded models with
+those fields.
+
+Clashing field names
+~~~~~~~~~~~~~~~~~~~~
+
+As with :class:`~.fields.PolymorphicEmbeddedModelField`, take care that your embedded
+models don't use :ref:`clashing field names
+<polymorphic-embedded-model-field-clashing-field-names>`.

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