Improve integration

Author: pfouque

.gitignore

@@ -133,3 +133,6 @@ test.db
 
 # django fsm command tests
 exports/*
+
+# Codex
+.codex

CHANGELOG.rst

@@ -1,6 +1,12 @@
 Changelog
 =========
 
+Unreleased
+~~~~~~~~~~
+
+- Add typing
+
+
 django-fsm-2 4.1.0 2025-11-03
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

django_fsm/__init__.py

@@ -97,7 +97,7 @@ class Transition:
     def __init__(
         self,
         method: Callable[..., _StateValue],
-        source: _StateValue | Sequence[_StateValue] | State,
+        source: _StateValue,
         target: _StateValue,
         on_error: _StateValue | None,
         conditions: list[_Condition] | None,
@@ -415,9 +415,7 @@ def get_all_transitions(self, instance_cls: type[_FSMModel]) -> Generator[Transi
         transitions = self.transitions[instance_cls]
 
         for transition in transitions.values():
-            meta = transition._django_fsm
-
-            yield from meta.transitions.values()
+            yield from transition._django_fsm.transitions.values()
 
     @override
     def contribute_to_class(
@@ -725,6 +723,8 @@ def has_transition_perm(bound_method: typing.Any, user: UserWithPermissions) ->
 
 
 class State:
+    allowed_states: Sequence[_StateValue]
+
     def get_state(
         self,
         model: _FSMModel,
@@ -737,8 +737,8 @@ def get_state(
 
 
 class RETURN_VALUE(State):  # noqa: N801
-    def __init__(self, *allowed_states: Sequence[_StateValue]) -> None:
-        self.allowed_states = allowed_states if allowed_states else None
+    def __init__(self, *allowed_states: _StateValue) -> None:
+        self.allowed_states = allowed_states or []
 
     def get_state(
         self,
@@ -762,7 +762,7 @@ def __init__(
         states: Sequence[_StateValue] | None = None,
     ) -> None:
         self.func = func
-        self.allowed_states = states
+        self.allowed_states = states or []
 
     def get_state(
         self,

django_fsm/management/commands/graph_transitions.py

@@ -33,7 +33,8 @@ def one_fsm_fields_data(
     model: type[models.Model], field_name: str
 ) -> tuple[FSMFieldMixin, type[models.Model]]:
     field = model._meta.get_field(field_name)
-    assert isinstance(field, FSMFieldMixin)
+    if not isinstance(field, FSMFieldMixin):
+        raise LookupError(f"{field_name} is not an FSMField")  # noqa: TRY004
     return (field, model)
 
 
@@ -61,9 +62,9 @@ def generate_dot(  # noqa: C901, PLR0912
     for field, model in fields_data:
         sources: set[tuple[(str, str)]] = set()
         targets: set[tuple[str, str]] = set()
-        edges: set[tuple[str, str, tuple[tuple[str, ...]]]] = set()
+        edges: set[tuple[str, str, tuple[tuple[str, str]]]] = set()
         any_targets: set[tuple[_StateValue, str]] = set()
-        any_except_targets: set[tuple[str, str]] = set()
+        any_except_targets: set[tuple[_StateValue, str]] = set()
 
         # dump nodes and edges
         for transition in field.get_all_transitions(model):
@@ -80,6 +81,7 @@ def generate_dot(  # noqa: C901, PLR0912
                 if isinstance(transition.source, GET_STATE | RETURN_VALUE)
                 else ((transition.source, node_name(field, transition.source)),)
             )
+
             for source, source_name in source_name_pair:
                 if transition.on_error:
                     on_error_name = node_name(field, transition.on_error)
@@ -92,16 +94,10 @@ def generate_dot(  # noqa: C901, PLR0912
                     elif transition.source == "+":
                         any_except_targets.add((target, transition.name))
                     else:
-                        add_transition(
-                            source,
-                            target,
-                            transition.name,
-                            source_name,
-                            field,
-                            sources,
-                            targets,
-                            edges,
-                        )
+                        target_name = node_name(field, target)
+                        sources.add((source_name, node_label(field, source)))
+                        targets.add((target_name, node_label(field, target)))
+                        edges.add((source_name, target_name, (("label", transition.name),)))
 
         targets.update(
             {
@@ -134,46 +130,23 @@ def generate_dot(  # noqa: C901, PLR0912
         final_states = targets - sources
         for name, label in final_states:
             subgraph.node(name, label=label, shape="doublecircle")
+
         for name, label in (sources | targets) - final_states:
             subgraph.node(name, label=label, shape="circle")
             # Adding initial state notation
             if field.default and label == field.default:
                 initial_name = node_name(field, "_initial")
                 subgraph.node(name=initial_name, label="", shape="point")
-                subgraph.edge(initial_name, name)
+                subgraph.edge(tail_name=initial_name, head_name=name)
+
         for source_name, target_name, attrs in edges:
-            subgraph.edge(source_name, target_name, **dict(attrs))
+            subgraph.edge(tail_name=source_name, head_name=target_name, **dict(attrs))
 
         result.subgraph(subgraph)
 
     return result
 
 
-def add_transition(
-    transition_source: _StateValue,
-    transition_target: _StateValue,
-    transition_name: str,
-    source_name: str,
-    field: FSMFieldMixin,
-    sources: set[tuple[str, str]],
-    targets: set[tuple[str, str]],
-    edges: set[tuple[str, str, tuple[tuple[str, str], ...]]],
-) -> None:
-    target_name = node_name(field, transition_target)
-    sources.add((source_name, node_label(field, transition_source)))
-    targets.add((target_name, node_label(field, transition_target)))
-    edges.add((source_name, target_name, (("label", transition_name),)))
-
-
-def get_graphviz_layouts() -> set[str] | Sequence[str]:
-    try:
-        import graphviz
-    except ModuleNotFoundError:
-        return {"sfdp", "circo", "twopi", "dot", "neato", "fdp", "osage", "patchwork"}
-    else:
-        return graphviz.ENGINES  # type: ignore[no-any-return]
-
-
 class Command(BaseCommand):
     help = "Creates a GraphViz dot file with transitions for selected fields"
 
@@ -192,7 +165,7 @@ def add_arguments(self, parser: ArgumentParser) -> None:
             action="store",
             dest="layout",
             default="dot",
-            help=f"Layout to be used by GraphViz for visualization: {get_graphviz_layouts()}.",
+            help=f"Layout to be used by GraphViz for visualization: {graphviz.ENGINES}.",
         )
         parser.add_argument(
             "--exclude",
@@ -204,13 +177,6 @@ def add_arguments(self, parser: ArgumentParser) -> None:
         )
         parser.add_argument("args", nargs="*", help=("[appname[.model[.field]]]"))
 
-    def render_output(self, graph: graphviz.Digraph, **options: typing.Any) -> None:
-        filename, graph_format = options["outputfile"].rsplit(".", 1)
-
-        graph.engine = options["layout"]
-        graph.format = graph_format
-        graph.render(filename)
-
     def handle(self, *args: str, **options: typing.Any) -> None:
         fields_data: list[tuple[FSMFieldMixin, type[models.Model]]] = []
         if args:
@@ -232,8 +198,11 @@ def handle(self, *args: str, **options: typing.Any) -> None:
 
         dotdata = generate_dot(fields_data, ignore_transitions=options["exclude"].split(","))
 
-        outputfile = options["outputfile"]
-        if outputfile:
-            self.render_output(dotdata, **options)
+        if outputfile := options["outputfile"]:
+            filename, graph_format = outputfile.rsplit(".", 1)
+
+            dotdata.engine = options["layout"]
+            dotdata.format = graph_format
+            dotdata.render(filename)
         else:
             self.stdout.write(str(dotdata))

pyproject.toml

@@ -25,11 +25,12 @@ classifiers = [
     "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: 3.13",
     "Programming Language :: Python :: 3.14",
+    "Typing :: Typed",
     "Topic :: Software Development :: Libraries :: Python Modules",
 ]
 dependencies = [
     "django>=4.2",
-    "typing-extensions>=4.13.2",
+    "typing-extensions; python_version < '3.12'",
 ]
 
 [project.urls]
@@ -46,7 +47,7 @@ graphviz = [
 dev = [
     "coverage",
     "django-guardian",
-    "graphviz",
+    "django-stubs-ext",
     "pre-commit",
     "pytest",
     "pytest-cov",
@@ -57,6 +58,7 @@ dev = [
 [tool.uv]
 default-groups = [
     "dev",
+    "graphviz"
 ]
 
 [tool.hatch.build.targets.sdist]
@@ -156,9 +158,3 @@ module = [
     "django_fsm.tests.*"
 ]
 disallow_untyped_defs = false
-
-[[tool.mypy.overrides]]
-module = [
-    "django_fsm.management.commands.graph_transitions.*",
-]
-ignore_errors = true

tests/settings.py

@@ -12,8 +12,14 @@
 
 from __future__ import annotations
 
+import typing
 from pathlib import Path
 
+if typing.TYPE_CHECKING:
+    import django_stubs_ext
+
+    django_stubs_ext.monkeypatch()
+
 # Build paths inside the project like this: BASE_DIR / 'subdir'.
 BASE_DIR = Path(__file__).resolve().parent.parent
 

tests/testapp/models.py

@@ -22,7 +22,7 @@ class Application(models.Model):
 
     state = FSMField(default="new")
 
-    @transition(field=state, source="new", target="published")
+    @transition(field=state, source="new", target="published", on_error="failed")
     def standard(self) -> None:
         pass