Add type annotations

setup.py

@@ -44,4 +44,5 @@ def read_relative_file(filename):
         "Programming Language :: Python :: 3.12",
         "Topic :: Utilities",
     ],
+    package_data={"typedmodels": ["py.typed"]},
 )

typedmodels/admin.py

@@ -1,15 +1,35 @@
-from django.contrib import admin
-from .models import TypedModel
+from typing import TYPE_CHECKING, Optional, Sequence, Callable, Any, Generic, Type
 
-class TypedModelAdmin(admin.ModelAdmin):
-    def get_fields(self, request, obj=None):
+from django.contrib.admin import ModelAdmin
+
+from .models import TypedModel, TypedModelT
+
+if TYPE_CHECKING:
+    from django.http import HttpRequest
+    from django.forms.forms import BaseForm
+
+
+class TypedModelAdmin(ModelAdmin, Generic[TypedModelT]):
+    model: "Type[TypedModelT]"
+
+    def get_fields(
+        self,
+        request: "HttpRequest",
+        obj: "Optional[TypedModelT]" = None,
+    ) -> Sequence[Callable[..., Any] | str]:
         fields = super().get_fields(request, obj)
         # we remove the type field from the admin of subclasses.
         if TypedModel not in self.model.__bases__:
             fields.remove(self.model._meta.get_field('type').name)
         return fields
 
-    def save_model(self, request, obj, form, change):
+    def save_model(
+        self,
+        request: "HttpRequest",
+        obj: "TypedModelT",
+        form: "BaseForm",
+        change,
+    ) -> None:
         if getattr(obj, '_typedmodels_type', None) is None:
             # new instances don't have the type attribute
             obj._typedmodels_type = form.cleaned_data['type']

typedmodels/models.py

@@ -1,3 +1,5 @@
+from typing import TYPE_CHECKING, Optional, Type, TypeVar, Generic
+
 from functools import partial
 import types
 
@@ -11,13 +13,21 @@
 from django.db.models.options import make_immutable_fields_list
 from django.utils.encoding import smart_str
 
+if TYPE_CHECKING:
+    from django.db.models import Model, QuerySet
+
+
+TypedModelT = TypeVar("TypedModelT", bound="TypedModel")
+
 
-class TypedModelManager(models.Manager):
-    def get_queryset(self):
+class TypedModelManager(models.Manager, Generic[TypedModelT]):
+    model: "Type[TypedModelT]"
+
+    def get_queryset(self) -> "QuerySet[TypedModelT]":
         qs = super(TypedModelManager, self).get_queryset()
         return self._filter_by_type(qs)
 
-    def _filter_by_type(self, qs):
+    def _filter_by_type(self, qs: "QuerySet[TypedModelT]"):
         if hasattr(self.model, "_typedmodels_type"):
             if len(self.model._typedmodels_subtypes) > 1:
                 qs = qs.filter(type__in=self.model._typedmodels_subtypes)
@@ -201,7 +211,7 @@ def do_related_class(self, other, cls):
 
             # add a get_type_classes classmethod to allow fetching of all the subclasses (useful for admin)
 
-            def get_type_classes(subcls):
+            def get_type_classes(subcls: "Type[TypedModelT]"):
                 if subcls is cls:
                     return list(cls._typedmodels_registry.values())
                 else:
@@ -212,7 +222,7 @@ def get_type_classes(subcls):
 
             cls.get_type_classes = classmethod(get_type_classes)
 
-            def get_types(subcls):
+            def get_types(subcls: "Type[TypedModelT]"):
                 if subcls is cls:
                     return list(cls._typedmodels_registry.keys())
                 else:
@@ -223,7 +233,7 @@ def get_types(subcls):
         return cls
 
     @staticmethod
-    def _model_has_field(cls, base_class, field_name):
+    def _model_has_field(cls, base_class: "Type[TypedModelT]", field_name: str):
         if field_name in base_class._meta._typedmodels_original_many_to_many:
             return True
         if field_name in base_class._meta._typedmodels_original_fields:
@@ -242,7 +252,7 @@ def _model_has_field(cls, base_class, field_name):
         return False
 
     @staticmethod
-    def _patch_fields_cache(cls, base_class):
+    def _patch_fields_cache(cls, base_class: "Type[TypedModelT]"):
         orig_get_fields = cls._meta._get_fields
 
         if django.VERSION >= (5, 0):
@@ -358,6 +368,8 @@ class Feline(Animal):
             def say_something(self):
                 return "meoww"
     '''
+    _typedmodels_type: Optional[str]
+    _typedmodels_subtypes: Optional[list[str]]
 
     objects = TypedModelManager()
 
@@ -406,7 +418,7 @@ def from_db(cls, db, field_names, values):
         new._state.db = db
         return new
 
-    def __init__(self, *args, _typedmodels_do_recast=None, **kwargs):
+    def __init__(self, *args, _typedmodels_do_recast=None, **kwargs) -> None:
         # Calling __init__ on base class because some functions (e.g. save()) need access to field values from base
         # class.
 
@@ -433,7 +445,7 @@ def __init__(self, *args, _typedmodels_do_recast=None, **kwargs):
         if _typedmodels_do_recast:
             self.recast()
 
-    def recast(self, typ=None):
+    def recast(self, typ: "Optional[Type[TypedModel]]" = None) -> None:
         for base in reversed(self.__class__.mro()):
             if issubclass(base, TypedModel) and hasattr(base, "_typedmodels_registry"):
                 break
@@ -471,10 +483,14 @@ def recast(self, typ=None):
         if current_cls != correct_cls:
             self.__class__ = correct_cls
 
-    def save(self, *args, **kwargs):
+    def save(self, *args, **kwargs) -> None:
+        self.presave(*args, **kwargs)
+        return super(TypedModel, self).save(*args, **kwargs)
+
+    def presave(self, *args, **kwargs) -> None:
+        """Perform checks before saving the model."""
         if not getattr(self, "_typedmodels_type", None):
             raise RuntimeError("Untyped %s cannot be saved." % self.__class__.__name__)
-        return super(TypedModel, self).save(*args, **kwargs)
 
     def _get_unique_checks(self, exclude=None, **kwargs):
         unique_checks, date_checks = super(TypedModel, self)._get_unique_checks(
@@ -498,7 +514,7 @@ def _get_unique_checks(self, exclude=None, **kwargs):
 _python_serializer_get_dump_object = _PythonSerializer.get_dump_object
 
 
-def _get_dump_object(self, obj):
+def _get_dump_object(self, obj: "Model") -> dict:
     if isinstance(obj, TypedModel):
         return {
             "pk": smart_str(obj._get_pk_val(), strings_only=True),
@@ -514,7 +530,7 @@ def _get_dump_object(self, obj):
 _xml_serializer_start_object = _XmlSerializer.start_object
 
 
-def _start_object(self, obj):
+def _start_object(self, obj: "Model") -> None:
     if isinstance(obj, TypedModel):
         self.indent(1)
         obj_pk = obj._get_pk_val()