refactor: enrich diagram IR to separate extraction from rendering

Move domain analysis logic from the renderer (dot.py) into the
extractor (extract.py) so the renderer becomes a pure IR→pydot mapping:

- Add ActionType enum replacing free strings for DiagramAction.type
- Add compound_state_ids and bidirectional_compound_ids to DiagramGraph
- Add DiagramTransition.is_initial flag for implicit initial transitions
- Remove redundant DiagramTransition.target (use targets list)
- Move _collect_compound_ids, _collect_compound_bidir_ids from renderer
  to extractor
- Add _mark_initial_transitions and _resolve_initial_states in extractor
- Remove _is_initial_candidate from renderer (use state.is_initial)
- Remove implicit transition filtering logic from renderer (use
  transition.is_initial)

Author: fgmacedo

statemachine/contrib/diagram/extract.py

@@ -1,7 +1,9 @@
 from typing import TYPE_CHECKING
 from typing import List
+from typing import Set
 from typing import Union
 
+from .model import ActionType
 from .model import DiagramAction
 from .model import DiagramGraph
 from .model import DiagramState
@@ -54,16 +56,16 @@ def _extract_state_actions(state: "State", getter) -> List[DiagramAction]:
     exit_ = str(getter(state.exit))
 
     if entry:
-        actions.append(DiagramAction(type="entry", body=entry))
+        actions.append(DiagramAction(type=ActionType.ENTRY, body=entry))
     if exit_:
-        actions.append(DiagramAction(type="exit", body=exit_))
+        actions.append(DiagramAction(type=ActionType.EXIT, body=exit_))
 
     for transition in state.transitions:
         if transition.internal:
             on_text = str(getter(transition.on))
             if on_text:
                 actions.append(
-                    DiagramAction(type="internal", body=f"{transition.event} / {on_text}")
+                    DiagramAction(type=ActionType.INTERNAL, body=f"{transition.event} / {on_text}")
                 )
 
     return actions
@@ -105,14 +107,12 @@ def _extract_transitions_from_state(state: "State", getter) -> List[DiagramTrans
     for transition in state.transitions:
         targets = transition.targets if transition.targets else []
         target_ids = [t.id for t in targets]
-        primary_target = target_ids[0] if target_ids else None
 
         cond_strs = [str(c) for c in transition.cond]
 
         result.append(
             DiagramTransition(
                 source=transition.source.id,
-                target=primary_target,
                 targets=target_ids,
                 event=transition.event,
                 guards=cond_strs,
@@ -136,6 +136,79 @@ def _extract_all_transitions(states, getter) -> List[DiagramTransition]:
     return result
 
 
+def _collect_compound_ids(states: List[DiagramState]) -> Set[str]:
+    """Collect IDs of states that have children (compound/parallel)."""
+    result: Set[str] = set()
+    for state in states:
+        if state.children:
+            result.add(state.id)
+        result.update(_collect_compound_ids(state.children))
+    return result
+
+
+def _collect_bidirectional_compound_ids(
+    transitions: List[DiagramTransition],
+    compound_ids: Set[str],
+) -> Set[str]:
+    """Find compound states that have both outgoing and incoming explicit edges."""
+    outgoing: Set[str] = set()
+    incoming: Set[str] = set()
+    for t in transitions:
+        if t.is_internal:
+            continue
+        # Skip implicit initial transitions
+        if t.source in compound_ids and not t.event and t.targets:
+            continue
+        if t.source in compound_ids:
+            outgoing.add(t.source)
+        for target_id in t.targets:
+            if target_id in compound_ids:
+                incoming.add(target_id)
+    return outgoing & incoming
+
+
+def _mark_initial_transitions(
+    transitions: List[DiagramTransition],
+    compound_ids: Set[str],
+) -> None:
+    """Mark implicit initial transitions (compound state → child, no event)."""
+    for t in transitions:
+        if t.source in compound_ids and not t.event and t.targets and not t.is_internal:
+            t.is_initial = True
+
+
+def _resolve_initial_states(states: List[DiagramState]) -> None:
+    """Ensure exactly one state per level has is_initial=True.
+
+    Skips parallel areas and history states. Falls back to document order
+    (first non-history, non-parallel-area state) when no explicit initial exists.
+    Recurses into children.
+
+    Parallel areas (children of a parallel state) have their is_initial flag
+    cleared: all regions are auto-activated, so no initial arrow is needed.
+    """
+    # Clear is_initial on parallel areas — all children of a parallel state
+    # are simultaneously active; initial arrows would be misleading.
+    for s in states:
+        if s.is_parallel_area:
+            s.is_initial = False
+
+    candidates = [
+        s
+        for s in states
+        if s.type not in (StateType.HISTORY_SHALLOW, StateType.HISTORY_DEEP)
+        and not s.is_parallel_area
+    ]
+
+    has_explicit_initial = any(s.is_initial for s in candidates)
+    if not has_explicit_initial and candidates:
+        candidates[0].is_initial = True
+
+    for state in states:
+        if state.children:
+            _resolve_initial_states(state.children)
+
+
 def extract(machine_or_class: "MachineRef") -> DiagramGraph:
     """Extract a DiagramGraph IR from a state machine instance or class.
 
@@ -171,8 +244,15 @@ class itself thanks to the metaclass.  Active-state highlighting is only
 
     transitions = _extract_all_transitions(machine.states, getter)
 
+    compound_ids = _collect_compound_ids(states)
+    bidir_ids = _collect_bidirectional_compound_ids(transitions, compound_ids)
+    _mark_initial_transitions(transitions, compound_ids)
+    _resolve_initial_states(states)
+
     return DiagramGraph(
         name=machine.name,
         states=states,
         transitions=transitions,
+        compound_state_ids=compound_ids,
+        bidirectional_compound_ids=bidir_ids,
     )

statemachine/contrib/diagram/model.py

@@ -2,7 +2,7 @@
 from dataclasses import field
 from enum import Enum
 from typing import List
-from typing import Optional
+from typing import Set
 
 
 class StateType(Enum):
@@ -19,9 +19,15 @@ class StateType(Enum):
     TERMINATE = "terminate"
 
 
+class ActionType(Enum):
+    ENTRY = "entry"
+    EXIT = "exit"
+    INTERNAL = "internal"
+
+
 @dataclass
 class DiagramAction:
-    type: str  # "entry", "exit", "internal"
+    type: ActionType
     body: str
 
 
@@ -40,16 +46,18 @@ class DiagramState:
 @dataclass
 class DiagramTransition:
     source: str
-    target: Optional[str]
     targets: List[str] = field(default_factory=list)
     event: str = ""
     guards: List[str] = field(default_factory=list)
     actions: List[str] = field(default_factory=list)
     is_internal: bool = False
+    is_initial: bool = False
 
 
 @dataclass
 class DiagramGraph:
     name: str
     states: List[DiagramState] = field(default_factory=list)
     transitions: List[DiagramTransition] = field(default_factory=list)
+    compound_state_ids: Set[str] = field(default_factory=set)
+    bidirectional_compound_ids: Set[str] = field(default_factory=set)

statemachine/contrib/diagram/renderers/dot.py

@@ -7,6 +7,7 @@
 
 import pydot
 
+from ..model import ActionType
 from ..model import DiagramAction
 from ..model import DiagramGraph
 from ..model import DiagramState
@@ -50,39 +51,12 @@ def __init__(self, config: Optional[DotRendererConfig] = None):
 
     def render(self, graph: DiagramGraph) -> pydot.Dot:
         """Render a DiagramGraph to a pydot.Dot object."""
-        self._collect_compound_ids(graph.states)
-        self._compound_bidir_ids = self._collect_compound_bidir_ids(graph.transitions)
+        self._compound_ids = graph.compound_state_ids
+        self._compound_bidir_ids = graph.bidirectional_compound_ids
         dot = self._create_graph(graph.name)
         self._render_states(graph.states, graph.transitions, dot)
         return dot
 
-    def _collect_compound_bidir_ids(self, transitions: List[DiagramTransition]) -> Set[str]:
-        """Find compound states that have both outgoing and incoming explicit edges.
-
-        Returns the set of compound state IDs that participate in at least one
-        bidirectional pair, so we can give them separate in/out anchor nodes.
-        """
-        outgoing: Set[str] = set()
-        incoming: Set[str] = set()
-        for t in transitions:
-            if t.is_internal:
-                continue
-            if t.source in self._compound_ids and not t.event and t.targets:
-                continue
-            if t.source in self._compound_ids:
-                outgoing.add(t.source)
-            for target_id in t.targets:
-                if target_id in self._compound_ids:
-                    incoming.add(target_id)
-        return outgoing & incoming
-
-    def _collect_compound_ids(self, states: List[DiagramState]) -> None:
-        """Pre-collect IDs of states that have children (compound/parallel)."""
-        for state in states:
-            if state.children:
-                self._compound_ids.add(state.id)
-            self._collect_compound_ids(state.children)
-
     def _create_graph(self, name: str) -> pydot.Dot:
         cfg = self.config
         graph_attrs = {
@@ -151,7 +125,7 @@ def _render_states(
         extra_nodes: Optional[List[pydot.Node]] = None,
     ) -> None:
         """Render states and transitions into the parent graph."""
-        initial_state = next((s for s in states if self._is_initial_candidate(s, states)), None)
+        initial_state = next((s for s in states if s.is_initial), None)
 
         # The atomic subgraph groups all non-compound states and the inner
         # initial dot (when inside a compound cluster) so Graphviz places them
@@ -255,22 +229,6 @@ def _render_initial_arrow(
         )
         return added_to_atomic
 
-    def _is_initial_candidate(self, state: DiagramState, siblings: List[DiagramState]) -> bool:
-        """Check if this state should get an initial arrow."""
-        # History states don't get initial arrows at this level
-        if state.type in (StateType.HISTORY_SHALLOW, StateType.HISTORY_DEEP):
-            return False
-        # All children of a parallel state are auto-initial; skip initial arrows
-        if state.is_parallel_area:
-            return False
-        # Use the is_initial flag from the model; fall back to document order
-        if state.is_initial:
-            return True
-        has_explicit_initial = any(s.is_initial for s in siblings)
-        if has_explicit_initial:
-            return False
-        return state is siblings[0] if siblings else False
-
     def _create_initial_node(self, node_id: str) -> pydot.Node:
         return pydot.Node(
             node_id,
@@ -293,7 +251,7 @@ def _create_atomic_node(self, state: DiagramState) -> pydot.Node:
         entry/exit actions embed an HTML TABLE (``border="0"``) inside the native
         shape to render UML-style compartments (name + separator + actions).
         """
-        actions = [a for a in state.actions if a.type != "internal" or a.body]
+        actions = [a for a in state.actions if a.type != ActionType.INTERNAL or a.body]
         fillcolor = self.config.state_active_fillcolor if state.is_active else "white"
         penwidth = self.config.state_active_penwidth if state.is_active else 2
 
@@ -363,9 +321,9 @@ def _build_html_table_label(
 
     @staticmethod
     def _format_action(action: DiagramAction) -> str:
-        if action.type == "internal":
+        if action.type == ActionType.INTERNAL:
             return action.body
-        return f"{action.type} / {action.body}"
+        return f"{action.type.value} / {action.body}"
 
     def _create_history_node(self, state: DiagramState) -> pydot.Node:
         label = "H*" if state.type == StateType.HISTORY_DEEP else "H"
@@ -443,7 +401,7 @@ def _build_compound_label(self, state: DiagramState) -> str:
         if state.type == StateType.PARALLEL:
             return f"<b>{name}</b> &#9783;"
 
-        actions = [a for a in state.actions if a.type != "internal" or a.body]
+        actions = [a for a in state.actions if a.type != ActionType.INTERNAL or a.body]
         if not actions:
             return f"<b>{name}</b>"
 
@@ -465,13 +423,8 @@ def _add_transitions_for_state(
         for transition in all_transitions:
             if transition.source != state.id or transition.is_internal:
                 continue
-            # Skip implicit initial transitions from a compound/parallel state to its
-            # initial child — these are already represented by the black-dot initial node.
-            if (
-                transition.source in self._compound_ids
-                and not transition.event
-                and transition.targets
-            ):
+            # Skip implicit initial transitions — represented by the black-dot initial node.
+            if transition.is_initial:
                 continue
             for edge in self._create_edges(transition):
                 graph.add_edge(edge)

tests/test_contrib_diagram.py

@@ -318,7 +318,7 @@ def test_history_state_default_transition():
     """History state's default transition appears as an edge in the diagram."""
     from statemachine.contrib.diagram.model import DiagramTransition
 
-    transition = DiagramTransition(source="hist", target="child1", targets=["child1"], event="")
+    transition = DiagramTransition(source="hist", targets=["child1"], event="")
     renderer = DotRenderer()
     renderer._compound_ids = set()
     edges = renderer._create_edges(transition)
@@ -361,9 +361,7 @@ def test_multi_target_transition_diagram():
     """Edges are created for all targets of a multi-target transition."""
     from statemachine.contrib.diagram.model import DiagramTransition
 
-    transition = DiagramTransition(
-        source="source", target="target1", targets=["target1", "target2"], event="go"
-    )
+    transition = DiagramTransition(source="source", targets=["target1", "target2"], event="go")
     renderer = DotRenderer()
     renderer._compound_ids = set()
     edges = renderer._create_edges(transition)