WIP: Annotate test helpers

Author: mthuurne

tests/managers.py

@@ -1,15 +1,22 @@
 from __future__ import annotations
 
-from model_utils.managers import SoftDeletableManager, SoftDeletableQuerySet
+from model_utils.managers import (
+    ModelT,
+    QuerySet,
+    SoftDeletableManager,
+    SoftDeletableQuerySet,
+)
 
 
-class CustomSoftDeleteQuerySet(SoftDeletableQuerySet):
-    def only_read(self):
+class CustomSoftDeleteQuerySet(SoftDeletableQuerySet[ModelT]):
+    def only_read(self) -> QuerySet[ModelT]:
         return self.filter(is_read=True)
 
 
-class CustomSoftDeleteManager(SoftDeletableManager):
+class CustomSoftDeleteManager(SoftDeletableManager[ModelT]):
     _queryset_class = CustomSoftDeleteQuerySet
 
-    def only_read(self):
-        return self.get_queryset().only_read()
+    def only_read(self) -> QuerySet[ModelT]:
+        qs = self.get_queryset()
+        assert isinstance(qs, self._queryset_class), qs
+        return qs.only_read()

tests/models.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import ClassVar, TypeVar
+from typing import Any, ClassVar, TypeVar
 
 from django.db import models
 from django.db.models import Manager
@@ -40,7 +40,7 @@ class InheritanceManagerTestParent(models.Model):
         on_delete=models.CASCADE)
     objects: ClassVar[InheritanceManager[InheritanceManagerTestParent]] = InheritanceManager()
 
-    def __str__(self):
+    def __str__(self) -> str:
         return "{}({})".format(
             self.__class__.__name__[len('InheritanceManagerTest'):],
             self.pk,
@@ -213,7 +213,7 @@ class Tracked(models.Model):
 
     tracker = FieldTracker()
 
-    def save(self, *args, **kwargs):
+    def save(self, *args: Any, **kwargs: Any) -> None:
         """ No-op save() to ensure that FieldTracker.patch_save() works. """
         super().save(*args, **kwargs)
 

tests/test_fields/test_field_tracker.py

@@ -33,7 +33,7 @@ class FieldTrackerTestCase(TestCase):
 
     tracker = None
 
-    def assertHasChanged(self, **kwargs):
+    def assertHasChanged(self, **kwargs) -> None:
         tracker = kwargs.pop('tracker', self.tracker)
         for field, value in kwargs.items():
             if value is None:
@@ -42,20 +42,20 @@ def assertHasChanged(self, **kwargs):
             else:
                 self.assertEqual(tracker.has_changed(field), value)
 
-    def assertPrevious(self, **kwargs):
+    def assertPrevious(self, **kwargs) -> None:
         tracker = kwargs.pop('tracker', self.tracker)
         for field, value in kwargs.items():
             self.assertEqual(tracker.previous(field), value)
 
-    def assertChanged(self, **kwargs):
+    def assertChanged(self, **kwargs) -> None:
         tracker = kwargs.pop('tracker', self.tracker)
         self.assertEqual(tracker.changed(), kwargs)
 
-    def assertCurrent(self, **kwargs):
+    def assertCurrent(self, **kwargs) -> None:
         tracker = kwargs.pop('tracker', self.tracker)
         self.assertEqual(tracker.current(), kwargs)
 
-    def update_instance(self, **kwargs):
+    def update_instance(self, **kwargs) -> None:
         for field, value in kwargs.items():
             setattr(self.instance, field, value)
         self.instance.save()
@@ -885,11 +885,11 @@ def setUp(self) -> None:
         self.instance = Tracked.objects.create(number=1)
         self.tracker = self.instance.tracker
 
-    def assertChanged(self, *fields):
+    def assertChanged(self, *fields) -> None:
         for f in fields:
             self.assertTrue(self.tracker.has_changed(f))
 
-    def assertNotChanged(self, *fields):
+    def assertNotChanged(self, *fields) -> None:
         for f in fields:
             self.assertFalse(self.tracker.has_changed(f))
 
@@ -920,7 +920,7 @@ def test_context_manager_fields(self) -> None:
     def test_tracker_decorator(self) -> None:
 
         @Tracked.tracker
-        def tracked_method(obj):
+        def tracked_method(obj) -> None:
             obj.name = 'new'
             self.assertChanged('name')
 
@@ -931,7 +931,7 @@ def tracked_method(obj):
     def test_tracker_decorator_fields(self) -> None:
 
         @Tracked.tracker(fields=['name'])
-        def tracked_method(obj):
+        def tracked_method(obj) -> None:
             obj.name = 'new'
             obj.number += 1
             self.assertChanged('name', 'number')