feat: Restore FieldTracker behaviour around save rather than save_base

jazzband · Jul 5, 2021 · 78f19e8 · 78f19e8
1 parent 7339caf
commit 78f19e8
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 6 deletions.
diff --git a/model_utils/tracker.py b/model_utils/tracker.py
@@ -192,8 +192,9 @@ class FieldTracker:
 
     tracker_class = FieldInstanceTracker
 
-    def __init__(self, fields=None):
+    def __init__(self, fields=None, always_include_fields=None):
         self.fields = fields
+        self.always_include_fields = always_include_fields
 
     def get_field_map(self, cls):
         """Returns dict mapping fields names to model attribute names"""
@@ -232,23 +233,28 @@ def initialize_tracker(self, sender, instance, **kwargs):
         instance._instance_initialized = True
 
     def patch_save(self, model):
-        self._patch(model, 'save_base', 'update_fields')
+        self._patch(model, 'save', 'update_fields')
         self._patch(model, 'refresh_from_db', 'fields')
 
     def _patch(self, model, method, fields_kwarg):
         original = getattr(model, method)
 
         @wraps(original)
         def inner(instance, *args, **kwargs):
-            ret = original(instance, *args, **kwargs)
             update_fields = kwargs.get(fields_kwarg)
             if not update_fields:
                 fields = update_fields
             else:
-                fields = (
+                fields = {
                     field for field in update_fields if
                     field in self.fields
-                )
+                }
+                if self.always_include_fields:
+                    fields.update({
+                        field for field in self.always_include_fields if
+                        field in self.fields
+                    })
+            ret = original(instance, *args, **kwargs)
             getattr(instance, self.attname).set_saved_fields(
                 fields=fields
             )

diff --git a/tests/models.py b/tests/models.py
@@ -243,7 +243,7 @@ class TrackerTimeStamped(TimeStampedModel):
     number = models.IntegerField()
     mutable = MutableField(default=None)
 
-    tracker = FieldTracker()
+    tracker = FieldTracker(always_include_fields=['modified'])
 
 
 class TrackedFK(models.Model):