From 2269f810c8017a5a26716a075513f2fd39ba64b3 Mon Sep 17 00:00:00 2001 From: Hasan Ramezani Date: Tue, 1 Dec 2020 23:10:36 +0100 Subject: [PATCH] Fix TypeError in save when model inherits from both TimeStampModel and StatusModel. (Fixes GH-465) --- CHANGES.rst | 2 ++ model_utils/models.py | 16 ++++++---------- tests/models.py | 10 ++++++++++ tests/test_models/test_timestamped_model.py | 11 ++++++++++- 4 files changed, 28 insertions(+), 11 deletions(-) diff --git a/CHANGES.rst b/CHANGES.rst index 184a644..8657423 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -4,6 +4,8 @@ CHANGES Next release (tbc) ------------------ - Applied `isort` to codebase (Refs GH-#402) +- Fix `TypeError` in save when model inherits from both TimeStampModel + and StatusModel. (Fixes GH-465) 4.1.0 (2020-11-29) ------------------ diff --git a/model_utils/models.py b/model_utils/models.py index 95320b0..268db8c 100644 --- a/model_utils/models.py +++ b/model_utils/models.py @@ -32,10 +32,8 @@ class TimeStampedModel(models.Model): a parameter to the update field argument. """ update_fields = kwargs.get('update_fields', None) - if update_fields is not None: - update_fields = set(update_fields) - if update_fields: - kwargs['update_fields'] = update_fields.union({'modified'}) + if update_fields: + kwargs['update_fields'] = set(update_fields).union({'modified'}) super().save(*args, **kwargs) @@ -74,12 +72,10 @@ class StatusModel(models.Model): status_changed field is updated even if it is not given as a parameter to the update field argument. """ - if ( - 'update_fields' in kwargs - and 'status' in kwargs['update_fields'] - and 'status_changed' not in kwargs['update_fields'] - ): - kwargs['update_fields'] += ['status_changed'] + update_fields = kwargs.get('update_fields', None) + if update_fields and 'status' in update_fields: + kwargs['update_fields'] = set(update_fields).union({'status_changed'}) + super().save(*args, **kwargs) class Meta: diff --git a/tests/models.py b/tests/models.py index 4323c4d..c789267 100644 --- a/tests/models.py +++ b/tests/models.py @@ -447,3 +447,13 @@ class CustomNotPrimaryUUIDModel(models.Model): class SaveSignalHandlingTestModel(SaveSignalHandlingModel): name = models.CharField(max_length=20) + + +class TimeStampWithStatusModel(TimeStampedModel, StatusModel): + STATUS = Choices( + ("active", _("active")), + ("deleted", _("deleted")), + ("on_hold", _("on hold")), + ) + + test_field = models.PositiveSmallIntegerField(default=0) diff --git a/tests/test_models/test_timestamped_model.py b/tests/test_models/test_timestamped_model.py index 10e295a..a303923 100644 --- a/tests/test_models/test_timestamped_model.py +++ b/tests/test_models/test_timestamped_model.py @@ -4,7 +4,7 @@ from django.test import TestCase from freezegun import freeze_time from parameterized import parameterized -from tests.models import TimeStamp +from tests.models import TimeStamp, TimeStampWithStatusModel class TimeStampedModelTests(TestCase): @@ -149,3 +149,12 @@ class TimeStampedModelTests(TestCase): t1.save(update_fields=None) self.assertEqual(t1.modified, datetime(2020, 1, 2)) + + def test_model_inherit_timestampmodel_and_statusmodel(self): + with freeze_time(datetime(2020, 1, 1)): + t1 = TimeStampWithStatusModel.objects.create() + + with freeze_time(datetime(2020, 1, 2)): + t1.save(update_fields=['test_field', 'status']) + + self.assertEqual(t1.modified, datetime(2020, 1, 2))