add EmbeddedModelField support to ArrayField

Author: timgraham

django_mongodb_backend/fields/array.py

@@ -126,10 +126,12 @@ def db_type(self, connection):
         return "array"
 
     def get_db_prep_value(self, value, connection, prepared=False):
+        from ..fields import EmbeddedModelField
+
         if isinstance(value, list | tuple):
-            # Workaround for https://code.djangoproject.com/ticket/35982
-            # (fixed in Django 5.2).
-            if isinstance(self.base_field, DecimalField):
+            # DecimalField here is a workaround for
+            #  https://code.djangoproject.com/ticket/35982 (fixed in Django 5.2).
+            if isinstance(self.base_field, (DecimalField | EmbeddedModelField)):
                 return [self.base_field.get_db_prep_save(i, connection) for i in value]
             return [self.base_field.get_db_prep_value(i, connection, prepared=False) for i in value]
         return value

tests/model_fields_/models.py

@@ -158,3 +158,20 @@ class D(EmbeddedModel):
 class E(EmbeddedModel):
     name = models.CharField(max_length=100)
     value = models.IntegerField()
+
+
+# ArrayField + EmbeddedModelField
+class Review(EmbeddedModel):
+    title = models.CharField(max_length=255)
+    rating = models.IntegerField()
+
+    def __str__(self):
+        return self.title
+
+
+class Movie(models.Model):
+    title = models.CharField(max_length=255)
+    reviews = ArrayField(EmbeddedModelField(Review), null=True)
+
+    def __str__(self):
+        return self.title

tests/model_fields_/test_embedded_model.py

@@ -17,7 +17,7 @@
 from django_mongodb_backend.fields import EmbeddedModelField
 from django_mongodb_backend.models import EmbeddedModel
 
-from .models import A, Address, Author, B, Book, C, D, Data, E, Holder, Library
+from .models import A, Address, Author, B, Book, C, D, Data, E, Holder, Library, Movie, Review
 from .utils import truncate_ms
 
 
@@ -83,6 +83,63 @@ def test_pre_save(self):
         self.assertGreater(obj.data.auto_now, auto_now_two)
 
 
+class EmbeddedArrayTests(TestCase):
+    def test_save_load(self):
+        reviews = [
+            Review(title="The best", rating=10),
+            Review(title="Mediocre", rating=5),
+            Review(title="Horrible", rating=1),
+        ]
+        Movie.objects.create(title="Lion King", reviews=reviews)
+        movie = Movie.objects.get(title="Lion King")
+        self.assertEqual(movie.reviews[0].title, "The best")
+        self.assertEqual(movie.reviews[0].rating, 10)
+        self.assertEqual(movie.reviews[1].title, "Mediocre")
+        self.assertEqual(movie.reviews[1].rating, 5)
+        self.assertEqual(movie.reviews[2].title, "Horrible")
+        self.assertEqual(movie.reviews[2].rating, 1)
+        self.assertEqual(len(movie.reviews), 3)
+
+    def test_save_load_null(self):
+        movie = Movie.objects.create(title="Lion King")
+        movie = Movie.objects.get(title="Lion King")
+        self.assertIsNone(movie.reviews)
+
+
+class EmbeddedArrayQueryingTests(TestCase):
+    @classmethod
+    def setUpTestData(cls):
+        reviews = [
+            Review(title="The best", rating=10),
+            Review(title="Mediocre", rating=5),
+            Review(title="Horrible", rating=1),
+        ]
+        cls.clouds = Movie.objects.create(title="Clouds", reviews=reviews)
+        reviews = [
+            Review(title="Super", rating=9),
+            Review(title="Meh", rating=5),
+            Review(title="Horrible", rating=2),
+        ]
+        cls.frozen = Movie.objects.create(title="Frozen", reviews=reviews)
+        reviews = [
+            Review(title="Excellent", rating=9),
+            Review(title="Wow", rating=8),
+            Review(title="Classic", rating=7),
+        ]
+        cls.bears = Movie.objects.create(title="Bears", reviews=reviews)
+
+    def test_filter_with_field(self):
+        self.assertCountEqual(
+            Movie.objects.filter(reviews__title="Horrible"), [self.clouds, self.frozen]
+        )
+
+    def test_filter_with_model(self):
+        self.assertCountEqual(
+            Movie.objects.filter(reviews=Review(title="Horrible", rating=2)),
+            [self.clouds, self.frozen],
+        )
+
+
 class QueryingTests(TestCase):
     @classmethod
     def setUpTestData(cls):