Handle excluding many-to-many fields with javascript_exclude.

Author: adamghill

django_unicorn/serializer.py

@@ -99,11 +99,17 @@ def _get_model_dict(model: Model) -> dict:
 
     for field in model._meta.get_fields():
         if field.is_relation and field.many_to_many:
+            exclude_field_attributes = getattr(
+                model, "__unicorn__exclude_field_attribute", []
+            )
             related_name = field.name
 
             if field.auto_created:
                 related_name = field.related_name or f"{field.name}_set"
 
+            if related_name in exclude_field_attributes:
+                continue
+
             pks = []
 
             try:
@@ -231,10 +237,6 @@ def _exclude_field_attributes(
     _exclude_field_attributes({"1": {"2": {"3": "4"}}}, ("1.2.3",)) == {"1": {"2": {}}}
     """
 
-    assert exclude_field_attributes is None or is_non_string_sequence(
-        exclude_field_attributes
-    ), "exclude_field_attributes type needs to be a sequence"
-
     if exclude_field_attributes:
         for field in exclude_field_attributes:
             field_splits = field.split(".")
@@ -278,6 +280,40 @@ def dumps(
     Returns a `str` instead of `bytes` (which deviates from `orjson.dumps`), but seems more useful.
     """
 
+    assert exclude_field_attributes is None or is_non_string_sequence(
+        exclude_field_attributes
+    ), "exclude_field_attributes type needs to be a sequence"
+
+    # This explicitly handles excluding many-to-many fields with a hacky private `__unicorn__exclude_field_attribute`
+    # attribute that gets used later in `_get_model_json`
+    if exclude_field_attributes:
+        updated_exclude_field_attributes = set()
+
+        for field_attributes in exclude_field_attributes:
+            field_attributes_split = field_attributes.split(".")
+            field_attribute = field_attributes_split[0]
+
+            for key in data.keys():
+                if isinstance(data[key], Model) and key == field_attribute:
+                    remaining_field_attributes = field_attributes_split[1:]
+
+                    if hasattr(data[key], "__unicorn__exclude_field_attribute"):
+                        data[key].__unicorn__exclude_field_attribute.append(
+                            remaining_field_attributes[0]
+                        )
+                    else:
+                        setattr(
+                            data[key],
+                            "__unicorn__exclude_field_attribute",
+                            remaining_field_attributes,
+                        )
+                    break
+                else:
+                    updated_exclude_field_attributes.add(field_attributes)
+
+        # Convert list to tuple again so it's hashable for `lru_cache`
+        exclude_field_attributes = tuple(updated_exclude_field_attributes)
+
     serialized_data = orjson.dumps(data, default=_json_serializer)
 
     if fix_floats:

tests/components/test_component.py

@@ -202,11 +202,14 @@ class TestComponent(UnicornView):
         another = {"Neutral Milk Hotel": {"album": {"On Avery Island": 1996}}}
 
         class Meta:
-            javascript_exclude = ("name.Universe", "another.Neutral Milk Hotel.album")
+            javascript_exclude = ("another.Neutral Milk Hotel.album",)
 
-    expected = '{"another":{"Neutral Milk Hotel":{}},"name":{}}'
+    expected = (
+        '{"another":{"Neutral Milk Hotel":{}},"name":{"Universe":{"World":"Earth"}}}'
+    )
     component = TestComponent(component_id="asdf1234", component_name="hello-world")
-    assert expected == component.get_frontend_context_variables()
+    actual = component.get_frontend_context_variables()
+    assert expected == actual
 
 
 def test_meta_javascript_exclude_nested_with_list():

tests/serializer/test_dumps.py

@@ -196,6 +196,90 @@ def test_model_foreign_key_recursive_parent():
     assert expected == actual
 
 
+@pytest.mark.django_db
+def test_model_many_to_many(django_assert_num_queries):
+    flavor_one = Flavor(name="name1", label="label1")
+    flavor_one.save()
+
+    taste1 = Taste(name="Bitter1")
+    taste1.save()
+    taste2 = Taste(name="Bitter2")
+    taste2.save()
+    taste3 = Taste(name="Bitter3")
+    taste3.save()
+
+    flavor_one.taste_set.add(taste1)
+    flavor_one.taste_set.add(taste2)
+    flavor_one.taste_set.add(taste3)
+
+    with django_assert_num_queries(2):
+        actual = serializer.dumps(flavor_one)
+
+    expected = {
+        "name": "name1",
+        "label": "label1",
+        "parent": None,
+        "float_value": None,
+        "decimal_value": None,
+        "uuid": str(flavor_one.uuid),
+        "datetime": None,
+        "date": None,
+        "time": None,
+        "duration": None,
+        "pk": 1,
+        "taste_set": [taste1.pk, taste2.pk, taste3.pk],
+        "origins": [],
+    }
+
+    assert expected == json.loads(actual)
+
+
+@pytest.mark.django_db
+def test_model_many_to_many_with_excludes(django_assert_num_queries):
+    flavor_one = Flavor(name="name1", label="label1")
+    flavor_one.save()
+
+    taste1 = Taste(name="Bitter1")
+    taste1.save()
+    taste2 = Taste(name="Bitter2")
+    taste2.save()
+    taste3 = Taste(name="Bitter3")
+    taste3.save()
+
+    flavor_one.taste_set.add(taste1)
+    flavor_one.taste_set.add(taste2)
+    flavor_one.taste_set.add(taste3)
+
+    # This shouldn't make any database calls because the many-to-manys are excluded and
+    # all of the other data is already set
+    with django_assert_num_queries(0):
+        actual = serializer.dumps(
+            {"flavor": flavor_one},
+            exclude_field_attributes=(
+                "flavor.taste_set",
+                "flavor.origins",
+            ),
+        )
+
+    expected = {
+        "flavor": {
+            "name": "name1",
+            "label": "label1",
+            "parent": None,
+            "float_value": None,
+            "decimal_value": None,
+            "uuid": str(flavor_one.uuid),
+            "datetime": None,
+            "date": None,
+            "time": None,
+            "duration": None,
+            "pk": 1,
+        }
+    }
+
+    assert expected == json.loads(actual)
+
+
 @pytest.mark.django_db
 def test_dumps_queryset(db):
     flavor_one = Flavor(name="name1", label="label1")
@@ -269,13 +353,62 @@ def test_get_model_dict():
 
 
 @pytest.mark.django_db
-def test_get_model_dict_many_to_many_is_referenced():
-    taste = Taste(name="Bitter")
-    taste.save()
+def test_get_model_dict_many_to_many_is_referenced(django_assert_num_queries):
     flavor_one = Flavor(name="name1", label="label1")
     flavor_one.save()
-    flavor_one.taste_set.add(taste)
-    actual = serializer._get_model_dict(flavor_one)
+
+    taste1 = Taste(name="Bitter")
+    taste1.save()
+    taste2 = Taste(name="Bitter2")
+    taste2.save()
+    taste3 = Taste(name="Bitter3")
+    taste3.save()
+
+    flavor_one.taste_set.add(taste1)
+    flavor_one.taste_set.add(taste2)
+    flavor_one.taste_set.add(taste3)
+
+    expected = {
+        "pk": 1,
+        "name": "name1",
+        "label": "label1",
+        "parent": None,
+        "decimal_value": None,
+        "float_value": None,
+        "uuid": str(flavor_one.uuid),
+        "date": None,
+        "datetime": None,
+        "time": None,
+        "duration": None,
+        "taste_set": [taste1.pk, taste2.pk, taste3.pk],
+        "origins": [],
+    }
+
+    flavor_one = Flavor.objects.filter(id=flavor_one.id).first()
+
+    with django_assert_num_queries(2):
+        actual = serializer._get_model_dict(flavor_one)
+
+    assert expected == actual
+
+
+@pytest.mark.django_db
+def test_get_model_dict_many_to_many_is_referenced_prefetched(
+    django_assert_num_queries,
+):
+    flavor_one = Flavor(name="name1", label="label1")
+    flavor_one.save()
+
+    taste1 = Taste(name="Bitter")
+    taste1.save()
+    taste2 = Taste(name="Bitter2")
+    taste2.save()
+    taste3 = Taste(name="Bitter3")
+    taste3.save()
+
+    flavor_one.taste_set.add(taste1)
+    flavor_one.taste_set.add(taste2)
+    flavor_one.taste_set.add(taste3)
 
     expected = {
         "pk": 1,
@@ -289,10 +422,18 @@ def test_get_model_dict_many_to_many_is_referenced():
         "datetime": None,
         "time": None,
         "duration": None,
-        "taste_set": [taste.pk],
+        "taste_set": [taste1.pk, taste2.pk, taste3.pk],
         "origins": [],
     }
 
+    flavor_one = (
+        Flavor.objects.prefetch_related("taste_set").filter(id=flavor_one.id).first()
+    )
+
+    # prefetch_related should reduce the database calls
+    with django_assert_num_queries(1):
+        actual = serializer._get_model_dict(flavor_one)
+
     assert expected == actual