Introducing IommiModel

iommi/form.py

@@ -332,10 +332,10 @@ def create_q_objects(suffix):
 
     when_clauses = [When(q, then=rank) for rank, q in enumerate(q_objects)]
     choices = field.choices.annotate(
-        iommi_ranking=Case(*when_clauses, default=len(q_objects) + 1, output_field=IntegerField())
+        _iommi_ranking=Case(*when_clauses, default=len(q_objects) + 1, output_field=IntegerField())
     )
 
-    return choices.filter(reduce(or_, q_objects)).order_by('iommi_ranking', *field.search_fields)
+    return choices.filter(reduce(or_, q_objects)).order_by('_iommi_ranking', *field.search_fields)
 
 
 def choice_queryset__parse(field, string_value, **_):

iommi/model.py

@@ -0,0 +1,50 @@
+from django.core.exceptions import FieldDoesNotExist
+from django.db.models import Model
+
+
+class IommiModel(Model):
+    class Meta:
+        abstract = True
+
+    iommi_ignored_attributes = ()
+
+    def __setattr__(self, name, value):
+        if not name.startswith('_') and name != 'pk':
+            try:
+                field = self._meta.get_field(name)
+
+                if getattr(field, 'primary_key', False):
+                    return object.__setattr__(self, name, value)
+
+            except FieldDoesNotExist as e:
+                if name not in self.iommi_ignored_attributes:
+                    raise TypeError(
+                        f'There is no field {name} on the model {self.__class__.__name__}. '
+                        f'You can assign arbitrary attributes if they start with `_`. '
+                        f'If this is an annotation, please add a tuple on the class named `iommi_ignored_attributes`'
+                        f'of valid annotated attributes that should not trigger this message.'
+                    ) from e
+
+            self.get_updated_fields().add(name)
+
+        return object.__setattr__(self, name, value)
+
+    def get_updated_fields(self):
+        return self.__dict__.setdefault('_updated_fields', set())
+
+    @classmethod
+    def from_db(cls, db, field_names, values):
+        result = super().from_db(db, field_names, values)
+        result.get_updated_fields().clear()
+        return result
+
+    def save(self, force_insert=False, force_update=False, using=None, update_fields=None):
+        if self.pk is not None and not force_insert:
+            update_fields = self.get_updated_fields()
+
+        super().save(force_insert=force_insert, force_update=force_update, using=using, update_fields=update_fields)
+
+        self.get_updated_fields().clear()
+
+    def __repr__(self):
+        return f'<{self.__class__.__name__} pk={self.pk}>'

iommi/model__tests.py

@@ -0,0 +1,108 @@
+import pytest
+from django.db.models import (
+    AutoField,
+    CASCADE,
+    Max,
+    OneToOneField,
+)
+
+from iommi.model import IommiModel
+from tests.models import (
+    MyAnnotatedIommiModel,
+    MyIommiModel,
+    NoRaceConditionModel,
+    RaceConditionModel,
+)
+
+
+def test_model():
+    m = MyIommiModel(foo=17)
+    assert m.foo == 17
+
+    assert m.get_updated_fields() == {'foo'}
+
+
+def test_constructor_exception():
+    with pytest.raises(TypeError):
+        MyIommiModel(bar=17)
+
+
+def test_attribute_exception():
+    m = MyIommiModel()
+    m._bar = 17
+    with pytest.raises(TypeError):
+        m.bar = 17
+
+
+def test_reversed():
+    class MyOtherModel(IommiModel):
+        bar = OneToOneField(MyIommiModel, related_name='other', on_delete=CASCADE)
+
+    o = MyOtherModel(bar=MyIommiModel())
+    o.bar = MyIommiModel()
+
+    MyIommiModel().other = MyOtherModel()
+
+
+def test_updated_fields():
+    m = MyIommiModel()
+    m.foo = 17
+
+    assert m.get_updated_fields() == {'foo'}
+
+
+def test_ignore_pk_field():
+    class WeirdPKNameModel(IommiModel):
+        this_is_a_pk = AutoField(primary_key=True)
+
+    m = WeirdPKNameModel()
+
+    assert m.get_updated_fields() == set()
+
+
+@pytest.mark.django_db
+def test_race_condition_on_save():
+    m = RaceConditionModel.objects.create(a=1, b=2)
+    m2 = RaceConditionModel.objects.get(pk=m.pk)
+    m2.b = 7
+    m2.save()
+
+    m.a = 17
+    m.save()  # This save() overwrites the value of b
+    assert RaceConditionModel.objects.get(pk=m.pk).b == 2
+
+
+@pytest.mark.django_db
+def test_no_race_condition_on_save():
+    m = NoRaceConditionModel.objects.create(a=1, b=2)
+    m2 = NoRaceConditionModel.objects.get(pk=m.pk)
+    assert m2.get_updated_fields() == set()
+    m2.b = 7
+    assert m2.get_updated_fields() == {'b'}
+    m2.save()
+
+    m.a = 17
+    assert m.get_updated_fields() == {'a'}
+    m.save()  # This save() does NOT overwrite b!
+    assert NoRaceConditionModel.objects.get(pk=m.pk).b == 7
+    assert not m.get_updated_fields()
+
+
+@pytest.mark.django_db
+def test_annotation():
+    MyIommiModel.objects.create(foo=2)
+    with pytest.raises(TypeError):
+        MyIommiModel.objects.annotate(fisk=Max('foo')).get()
+
+    MyAnnotatedIommiModel.objects.create(foo=2)
+    MyAnnotatedIommiModel.objects.annotate(fisk=Max('foo')).get()
+
+
+@pytest.mark.django_db
+def test_force_insert():
+    MyIommiModel.objects.create(pk=3, foo=1)
+
+
+def test_repr():
+    assert repr(MyIommiModel(pk=None, foo=1)) == '<MyIommiModel pk=None>'
+    assert repr(MyIommiModel(pk=7, foo=1)) == '<MyIommiModel pk=7>'

tests/models.py

@@ -14,6 +14,7 @@
 )
 
 from iommi import register_search_fields
+from iommi.model import IommiModel
 
 
 class FormFromModelTest(Model):
@@ -153,8 +154,8 @@ class TBaz(Model):
 
     class Meta:
         ordering = ('pk',)
-        
-        
+
+
 class AdminUnique(Model):
     foo = IntegerField()
     unique = IntegerField(unique=True)
@@ -271,3 +272,23 @@ class CustomField(models.Field):
 
 class NotRegisteredCustomFieldModel(models.Model):
     custom_field = CustomField()
+
+
+class RaceConditionModel(models.Model):
+    a = IntegerField()
+    b = IntegerField()
+
+
+class NoRaceConditionModel(IommiModel):
+    a = IntegerField()
+    b = IntegerField()
+
+
+class MyIommiModel(IommiModel):
+    foo = IntegerField()
+
+
+class MyAnnotatedIommiModel(IommiModel):
+    foo = IntegerField()
+
+    iommi_ignored_attributes = ['fisk']