Step 3

Author: pfouque

.mypy.ini

@@ -15,10 +15,10 @@ check_untyped_defs = True
 # get passing if you use a lot of untyped libraries
 disallow_subclassing_any = True
 disallow_untyped_decorators = True
-; disallow_any_generics = True
+disallow_any_generics = True
 
 # These next few are various gradations of forcing use of type annotations
-; disallow_untyped_calls = True
+disallow_untyped_calls = True
 ; disallow_incomplete_defs = True
 ; disallow_untyped_defs = True
 

django_fsm/__init__.py

@@ -35,21 +35,30 @@
 
 if TYPE_CHECKING:
     from collections.abc import Callable
+    from collections.abc import Generator
     from collections.abc import Sequence
     from typing import Any
 
-    from django.contrib.auth.models import AbstractBaseUser
+    from django.contrib.auth.models import PermissionsMixin as UserWithPermissions
     from django.utils.functional import _StrOrPromise
 
     _Model = models.Model
+    _Field = models.Field[Any, Any]
+    CharField = models.CharField[str, str]
+    IntegerField = models.IntegerField[int, int]
+    ForeignKey = models.ForeignKey[Any, Any]
 else:
     _Model = object
+    _Field = object
+    CharField = models.CharField
+    IntegerField = models.IntegerField
+    ForeignKey = models.ForeignKey
 
 
 class TransitionNotAllowed(Exception):
     """Raised when a transition is not allowed"""
 
-    def __init__(self, *args, **kwargs) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         self.object = kwargs.pop("object", None)
         self.method = kwargs.pop("method", None)
         super().__init__(*args, **kwargs)
@@ -70,12 +79,12 @@ class ConcurrentTransition(Exception):
 class Transition:
     def __init__(
         self,
-        method: Callable,
+        method: Callable[..., Any],
         source: str | int | Sequence[str | int] | State,
         target: str | int | State | None,
         on_error: str | int | None,
         conditions: list[Callable[[Any], bool]],
-        permission: str | Callable[[models.Model, AbstractBaseUser], bool] | None,
+        permission: str | Callable[[models.Model, UserWithPermissions], bool] | None,
         custom: dict[str, _StrOrPromise],
     ) -> None:
         self.method = method
@@ -90,7 +99,7 @@ def __init__(
     def name(self) -> str:
         return self.method.__name__
 
-    def has_perm(self, instance, user) -> bool:
+    def has_perm(self, instance, user: UserWithPermissions) -> bool:
         if not self.permission:
             return True
         elif callable(self.permission):
@@ -114,7 +123,7 @@ def __eq__(self, other):
         return False
 
 
-def get_available_FIELD_transitions(instance, field):
+def get_available_FIELD_transitions(instance, field: FSMFieldMixin) -> Generator[Transition, None, None]:
     """
     List of transitions available in current model state
     with all conditions met
@@ -128,14 +137,16 @@ def get_available_FIELD_transitions(instance, field):
             yield meta.get_transition(curr_state)
 
 
-def get_all_FIELD_transitions(instance, field):
+def get_all_FIELD_transitions(instance, field: FSMFieldMixin) -> Generator[Transition, None, None]:
     """
     List of all transitions available in current model state
     """
     return field.get_all_transitions(instance.__class__)
 
 
-def get_available_user_FIELD_transitions(instance, user, field):
+def get_available_user_FIELD_transitions(
+    instance, user: UserWithPermissions, field: FSMFieldMixin
+) -> Generator[Transition, None, None]:
     """
     List of transitions available in current model state
     with all conditions met and user have rights on it
@@ -154,15 +165,24 @@ def __init__(self, field, method) -> None:
         self.field = field
         self.transitions: dict[str, Any] = {}  # source -> Transition
 
-    def get_transition(self, source):
+    def get_transition(self, source: str):
         transition = self.transitions.get(source, None)
         if transition is None:
             transition = self.transitions.get("*", None)
         if transition is None:
             transition = self.transitions.get("+", None)
         return transition
 
-    def add_transition(self, method, source, target, on_error=None, conditions=[], permission=None, custom={}) -> None:
+    def add_transition(
+        self,
+        method: Callable[..., Any],
+        source: str,
+        target: str | int,
+        on_error: str | int | None = None,
+        conditions: list[Callable[[Any], bool]] = [],
+        permission: str | Callable[[models.Model, UserWithPermissions], bool] | None = None,
+        custom: dict[str, _StrOrPromise] = {},
+    ) -> None:
         if source in self.transitions:
             raise AssertionError(f"Duplicate transition for {source} state")
 
@@ -204,7 +224,7 @@ def conditions_met(self, instance, state) -> bool:
         else:
             return all(map(lambda condition: condition(instance), transition.conditions))
 
-    def has_transition_perm(self, instance, state, user) -> bool:
+    def has_transition_perm(self, instance, state, user: UserWithPermissions) -> bool:
         transition = self.get_transition(state)
 
         if not transition:
@@ -247,10 +267,10 @@ def __set__(self, instance, value) -> None:
         self.field.set_state(instance, value)
 
 
-class FSMFieldMixin(Field):
+class FSMFieldMixin(_Field):
     descriptor_class = FSMFieldDescriptor
 
-    def __init__(self, *args, **kwargs) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         self.protected = kwargs.pop("protected", False)
         self.transitions: dict[Any, dict[str, Any]] = {}  # cls -> (transitions name -> method)
         self.state_proxy = {}  # state -> ProxyClsRef
@@ -275,15 +295,15 @@ def deconstruct(self):
             kwargs["protected"] = self.protected
         return name, path, args, kwargs
 
-    def get_state(self, instance):
+    def get_state(self, instance) -> Any:
         # The state field may be deferred. We delegate the logic of figuring this out
         # and loading the deferred field on-demand to Django's built-in DeferredAttribute class.
         return DeferredAttribute(self).__get__(instance)  # type: ignore[attr-defined]
 
-    def set_state(self, instance, state):
+    def set_state(self, instance, state: str) -> None:
         instance.__dict__[self.name] = state
 
-    def set_proxy(self, instance, state):
+    def set_proxy(self, instance, state: str) -> None:
         """
         Change class
         """
@@ -304,7 +324,7 @@ def set_proxy(self, instance, state):
 
             instance.__class__ = model
 
-    def change_state(self, instance, method, *args, **kwargs):
+    def change_state(self, instance, method, *args: Any, **kwargs: Any):
         meta = method._django_fsm
         method_name = method.__name__
         current_state = self.get_state(instance)
@@ -357,7 +377,7 @@ def change_state(self, instance, method, *args, **kwargs):
 
         return result
 
-    def get_all_transitions(self, instance_cls):
+    def get_all_transitions(self, instance_cls) -> Generator[Transition, None, None]:
         """
         Returns [(source, target, name, method)] for all field transitions
         """
@@ -384,7 +404,7 @@ def contribute_to_class(self, cls, name, private_only=False, **kwargs):
 
         class_prepared.connect(self._collect_transitions)
 
-    def _collect_transitions(self, *args, **kwargs):
+    def _collect_transitions(self, *args: Any, **kwargs: Any):
         sender = kwargs["sender"]
 
         if not issubclass(sender, self.base_cls):
@@ -413,25 +433,25 @@ def is_field_transition_method(attr):
         self.transitions[sender] = sender_transitions
 
 
-class FSMField(FSMFieldMixin, models.CharField):
+class FSMField(FSMFieldMixin, CharField):
     """
     State Machine support for Django model as CharField
     """
 
-    def __init__(self, *args, **kwargs) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         kwargs.setdefault("max_length", 50)
         super().__init__(*args, **kwargs)
 
 
-class FSMIntegerField(FSMFieldMixin, models.IntegerField):
+class FSMIntegerField(FSMFieldMixin, IntegerField):
     """
     Same as FSMField, but stores the state value in an IntegerField.
     """
 
     pass
 
 
-class FSMKeyField(FSMFieldMixin, models.ForeignKey):
+class FSMKeyField(FSMFieldMixin, ForeignKey):
     """
     State Machine support for Django model
     """
@@ -496,7 +516,7 @@ class ConcurrentTransitionMixin(_Model):
     state, thus practically negating their effect.
     """
 
-    def __init__(self, *args, **kwargs) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
         self._update_initial_state()
 
@@ -534,14 +554,14 @@ def _do_update(self, base_qs, using, pk_val, values, update_fields, forced_updat
 
         return updated
 
-    def _update_initial_state(self):
+    def _update_initial_state(self) -> None:
         self.__initial_states = {field.attname: field.value_from_object(self) for field in self.state_fields}
 
-    def refresh_from_db(self, *args, **kwargs):
+    def refresh_from_db(self, *args: Any, **kwargs: Any) -> None:
         super().refresh_from_db(*args, **kwargs)
         self._update_initial_state()
 
-    def save(self, *args, **kwargs):
+    def save(self, *args: Any, **kwargs: Any) -> None:
         super().save(*args, **kwargs)
         self._update_initial_state()
 
@@ -552,7 +572,7 @@ def transition(
     target: str | int | State | None = None,
     on_error: str | int | None = None,
     conditions: list[Callable[[Any], bool]] = [],
-    permission: str | Callable[[models.Model, AbstractBaseUser], bool] | None = None,
+    permission: str | Callable[[models.Model, UserWithPermissions], bool] | None = None,
     custom: dict[str, _StrOrPromise] = {},
 ):
     """
@@ -576,7 +596,7 @@ def inner_transition(func):
             func._django_fsm.add_transition(func, source, target, on_error, conditions, permission, custom)
 
         @wraps(func)
-        def _change_state(instance, *args, **kwargs):
+        def _change_state(instance, *args: Any, **kwargs: Any):
             return fsm_meta.field.change_state(instance, func, *args, **kwargs)
 
         if not wrapper_installed:
@@ -587,7 +607,7 @@ def _change_state(instance, *args, **kwargs):
     return inner_transition
 
 
-def can_proceed(bound_method, check_conditions=True) -> bool:
+def can_proceed(bound_method, check_conditions: bool = True) -> bool:
     """
     Returns True if model in state allows to call bound_method
 
@@ -604,7 +624,7 @@ def can_proceed(bound_method, check_conditions=True) -> bool:
     return meta.has_transition(current_state) and (not check_conditions or meta.conditions_met(self, current_state))
 
 
-def has_transition_perm(bound_method, user) -> bool:
+def has_transition_perm(bound_method, user: UserWithPermissions) -> bool:
     """
     Returns True if model in state allows to call bound_method and user have rights on it
     """
@@ -628,7 +648,7 @@ def get_state(self, model, transition, result, args=[], kwargs={}):
 
 
 class RETURN_VALUE(State):
-    def __init__(self, *allowed_states) -> None:
+    def __init__(self, *allowed_states: Sequence[str | int]) -> None:
         self.allowed_states = allowed_states if allowed_states else None
 
     def get_state(self, model, transition, result, args=[], kwargs={}):
@@ -639,7 +659,7 @@ def get_state(self, model, transition, result, args=[], kwargs={}):
 
 
 class GET_STATE(State):
-    def __init__(self, func, states=None) -> None:
+    def __init__(self, func: Callable[..., str | int], states: Sequence[str | int] | None = None) -> None:
         self.func = func
         self.allowed_states = states