Merge pull request #1198 from thunderstore-io/visibility-mixin-save-fix

Pass through update_fields and other kwargs in VisibilityMixin save override

Author: MythicManiac

django/thunderstore/community/models/community.py

@@ -19,7 +19,7 @@
 from thunderstore.core.exceptions import PermissionValidationError
 from thunderstore.core.mixins import TimestampMixin
 from thunderstore.core.types import UserType
-from thunderstore.core.utils import check_validity
+from thunderstore.core.utils import check_validity, extend_update_fields_if_present
 from thunderstore.permissions.utils import validate_user
 
 if TYPE_CHECKING:
@@ -134,67 +134,52 @@ def aggregated(self) -> "AggregatedFields":
             else CommunityAggregatedFields.get_empty()
         )
 
-    def save(self, *args, **kwargs):
+    def save(self, **kwargs):
         if self.pk:
             in_db = type(self).objects.get(pk=self.pk)
             if in_db.identifier != self.identifier:
                 raise ValidationError("Field 'identifier' is read only")
         if not self.icon:
             self.icon_width = 0
             self.icon_height = 0
-            if "update_fields" in kwargs:
-                kwargs["update_fields"] = set(
-                    kwargs["update_fields"]
-                    + (
-                        "icon_width",
-                        "icon_height",
-                    ),
-                )
+            kwargs = extend_update_fields_if_present(
+                kwargs,
+                "icon_width",
+                "icon_height",
+            )
         if not self.background_image:
             self.background_image_width = 0
             self.background_image_height = 0
-            if "update_fields" in kwargs:
-                kwargs["update_fields"] = set(
-                    kwargs["update_fields"]
-                    + (
-                        "background_image_width",
-                        "background_image_height",
-                    ),
-                )
+            kwargs = extend_update_fields_if_present(
+                kwargs,
+                "background_image_width",
+                "background_image_height",
+            )
         if not self.hero_image:
             self.hero_image_width = 0
             self.hero_image_height = 0
-            if "update_fields" in kwargs:
-                kwargs["update_fields"] = set(
-                    kwargs["update_fields"]
-                    + (
-                        "hero_image_width",
-                        "hero_image_height",
-                    ),
-                )
+            kwargs = extend_update_fields_if_present(
+                kwargs,
+                "hero_image_width",
+                "hero_image_height",
+            )
         if not self.cover_image:
             self.cover_image_width = 0
             self.cover_image_height = 0
-            if "update_fields" in kwargs:
-                kwargs["update_fields"] = set(
-                    kwargs["update_fields"]
-                    + (
-                        "cover_image_width",
-                        "cover_image_height",
-                    ),
-                )
+            kwargs = extend_update_fields_if_present(
+                kwargs,
+                "cover_image_width",
+                "cover_image_height",
+            )
         if not self.community_icon:
             self.community_icon_width = 0
             self.community_icon_height = 0
-            if "update_fields" in kwargs:
-                kwargs["update_fields"] = set(
-                    kwargs["update_fields"]
-                    + (
-                        "community_icon_width",
-                        "community_icon_height",
-                    ),
-                )
-        return super().save(*args, **kwargs)
+            kwargs = extend_update_fields_if_present(
+                kwargs,
+                "community_icon_width",
+                "community_icon_height",
+            )
+        return super().save(**kwargs)
 
     def __str__(self):
         return self.name

django/thunderstore/core/mixins.py

@@ -7,19 +7,16 @@
 from django.urls import reverse
 
 from thunderstore.cache.storage import CACHE_STORAGE
+from thunderstore.core.utils import extend_update_fields_if_present
 
 
 class TimestampMixin(models.Model):
     datetime_created = models.DateTimeField(auto_now_add=True)
     datetime_updated = models.DateTimeField(auto_now=True)
 
-    def save(self, *args, **kwargs):
-        update_fields = kwargs.pop("update_fields", [])
-        if update_fields:
-            kwargs["update_fields"] = tuple(
-                set(list(update_fields) + ["datetime_updated"])
-            )
-        super().save(*args, **kwargs)
+    def save(self, **kwargs):
+        kwargs = extend_update_fields_if_present(kwargs, "datetime_updated")
+        super().save(**kwargs)
 
     class Meta:
         abstract = True

django/thunderstore/core/tests/test_utils.py

@@ -7,6 +7,7 @@
 from thunderstore.core.utils import (
     capture_exception,
     check_validity,
+    extend_update_fields_if_present,
     make_full_url,
     replace_cdn,
     sanitize_filename,
@@ -187,3 +188,41 @@ def test_replace_cdn__raises_for_relative_urls(url: str) -> None:
         match="Absolute URL including protocol required",
     ):
         replace_cdn(url, "irrelevant")
+
+
+def test_extend_update_fields_if_present__no_update_fields_key() -> None:
+    original = {"some": "value"}
+    result = extend_update_fields_if_present(original, "new_field")
+    assert result is not original
+    # Should not add update_fields key when it's not present in original
+    assert "update_fields" not in result
+    # Original dict remains unchanged
+    assert original == {"some": "value"}
+
+
+def test_extend_update_fields_if_present__update_fields_none() -> None:
+    original = {"update_fields": None, "other": 1}
+    result = extend_update_fields_if_present(original, "x")
+    # Copy returned
+    assert result is not original
+    # When update_fields is None, it should be left as-is
+    assert result["update_fields"] is None
+    assert result["other"] == 1
+    # Original remains unmodified
+    assert original["update_fields"] is None
+
+
+def test_extend_update_fields_if_present__extends_and_deduplicates() -> None:
+    original = {"update_fields": ["a", "b"]}
+    result = extend_update_fields_if_present(original, "b", "c")
+    # Should return a set with union of existing and new fields
+    assert isinstance(result["update_fields"], set)
+    assert result["update_fields"] == {"a", "b", "c"}
+    # Original should not be mutated
+    assert original["update_fields"] == ["a", "b"]
+
+
+def test_extend_update_fields_if_present__accepts_tuple_and_empty_new_fields() -> None:
+    original = {"update_fields": ("x",)}
+    result = extend_update_fields_if_present(original)
+    assert result["update_fields"] == {"x"}

django/thunderstore/core/utils.py

@@ -1,6 +1,6 @@
 import re
 import urllib.parse
-from typing import Callable, Optional
+from typing import Any, Callable, Dict, List, Optional
 
 from django.conf import settings
 from django.core.exceptions import ValidationError
@@ -111,6 +111,16 @@ def sanitize_filepath(filepath: Optional[str]) -> Optional[str]:
     )
 
 
+def extend_update_fields_if_present(
+    original_kwargs: Dict[str, Any], *new_fields: str
+) -> Dict[str, Any]:
+    """Returns a copy of original_kwargs with the update_fields field updated"""
+    result = {**original_kwargs}
+    if (upfields := original_kwargs.get("update_fields")) is not None:
+        result["update_fields"] = {*upfields, *new_fields}
+    return result
+
+
 def validate_filepath_prefix(filepath: Optional[str]) -> Optional[str]:
     stripped = sanitize_filepath(filepath)
     if stripped != filepath:

django/thunderstore/permissions/mixins.py

@@ -5,6 +5,7 @@
 from django.db.models import Q
 
 from thunderstore.core.types import UserType
+from thunderstore.core.utils import extend_update_fields_if_present
 from thunderstore.permissions.models import VisibilityFlags
 
 
@@ -71,13 +72,14 @@ def set_zero_visibility(self):
         self.visibility.admin_list = False
 
     @transaction.atomic
-    def save(self, *args, **kwargs):
+    def save(self, **kwargs):
         if not self.pk and not self.visibility:
             self.visibility = VisibilityFlags.objects.create_public()
+            kwargs = extend_update_fields_if_present(kwargs, "visibility")
 
         self.update_visibility()
 
-        super().save()
+        super().save(**kwargs)
 
     class Meta:
         abstract = True