feat: Diagram compound and parallel states

Author: fgmacedo

conftest.py

@@ -1,4 +1,6 @@
+import io
 import sys
+from unittest import mock
 
 import pytest
 
@@ -31,3 +33,25 @@ def pytest_ignore_collect(collection_path, path, config):
 
     if "django_project" in str(path):
         return True
+
+
+@pytest.fixture(autouse=True, scope="module")
+def mock_dot_write(request):
+    def open_effect(
+        filename,
+        mode="r",
+        *args,
+        **kwargs,
+    ):
+        if mode in ("r", "rt", "rb"):
+            return open(filename, mode, *args, **kwargs)
+        elif filename.startswith("/tmp/"):
+            return open(filename, mode, *args, **kwargs)
+        elif "b" in mode:
+            return io.BytesIO()
+        else:
+            return io.StringIO()
+
+    with mock.patch("pydot.core.io.open", spec=True) as m:
+        m.side_effect = open_effect
+        yield m

docs/diagram.md

@@ -42,7 +42,7 @@ Graphviz. For example, on Debian-based systems (such as Ubuntu), you can use the
 >>> dot = graph()
 
 >>> dot.to_string()  # doctest: +ELLIPSIS
-'digraph list {...
+'digraph OrderControl {...
 
 ```
 

docs/images/order_control_machine_initial.png

@@ -32,21 +32,34 @@ class DotGraphMachine:
     def __init__(self, machine):
         self.machine = machine
 
-    def _get_graph(self):
-        machine = self.machine
+    def _get_graph(self, machine):
         return pydot.Dot(
-            "list",
+            machine.name,
             graph_type="digraph",
             label=machine.name,
             fontname=self.font_name,
             fontsize=self.state_font_size,
             rankdir=self.graph_rankdir,
+            compound="true",
         )
 
-    def _initial_node(self):
+    def _get_subgraph(self, state):
+        style = ", solid"
+        if state.parent and state.parent.parallel:
+            style = ", dashed"
+        subgraph = pydot.Subgraph(
+            label=f"{state.name}",
+            graph_name=f"cluster_{state.id}",
+            style=f"rounded{style}",
+            cluster="true",
+        )
+        return subgraph
+
+    def _initial_node(self, state):
         node = pydot.Node(
-            "i",
-            shape="circle",
+            self._state_id(state),
+            label="",
+            shape="point",
             style="filled",
             fontsize="1pt",
             fixedsize="true",
@@ -56,14 +69,18 @@ def _initial_node(self):
         node.set_fillcolor("black")
         return node
 
-    def _initial_edge(self):
+    def _initial_edge(self, initial_node, state):
+        extra_params = {}
+        if state.states:
+            extra_params["lhead"] = f"cluster_{state.id}"
         return pydot.Edge(
-            "i",
-            self.machine.initial_state.id,
+            initial_node.get_name(),
+            self._state_id(state),
             label="",
             color="blue",
             fontname=self.font_name,
             fontsize=self.transition_font_size,
+            **extra_params,
         )
 
     def _actions_getter(self):
@@ -104,11 +121,18 @@ def _state_actions(self, state):
 
         return actions
 
+    @staticmethod
+    def _state_id(state):
+        if state.states:
+            return f"{state.id}_anchor"
+        else:
+            return state.id
+
     def _state_as_node(self, state):
         actions = self._state_actions(state)
 
         node = pydot.Node(
-            state.id,
+            self._state_id(state),
             label=f"{state.name}{actions}",
             shape="rectangle",
             style="rounded, filled",
@@ -127,29 +151,64 @@ def _transition_as_edge(self, transition):
         cond = ", ".join([str(cond) for cond in transition.cond])
         if cond:
             cond = f"\n[{cond}]"
+
+        extra_params = {}
+        has_substates = transition.source.states or transition.target.states
+        if transition.source.states:
+            extra_params["ltail"] = f"cluster_{transition.source.id}"
+        if transition.target.states:
+            extra_params["lhead"] = f"cluster_{transition.target.id}"
+
         return pydot.Edge(
-            transition.source.id,
-            transition.target.id,
+            self._state_id(transition.source),
+            self._state_id(transition.target),
             label=f"{transition.event}{cond}",
             color="blue",
             fontname=self.font_name,
             fontsize=self.transition_font_size,
+            minlen=2 if has_substates else 1,
+            **extra_params,
         )
 
     def get_graph(self):
-        graph = self._get_graph()
-        graph.add_node(self._initial_node())
-        graph.add_edge(self._initial_edge())
+        graph = self._get_graph(self.machine)
+        self._graph_states(self.machine, graph)
+        return graph
 
-        for state in self.machine.states:
-            graph.add_node(self._state_as_node(state))
-            for transition in state.transitions:
+    def _graph_states(self, state, graph):
+        initial_node = self._initial_node(state)
+        initial_subgraph = pydot.Subgraph(
+            graph_name=f"{initial_node.get_name()}_initial",
+            label="",
+            peripheries=0,
+            margin=0,
+        )
+        atomic_states_subgraph = pydot.Subgraph(
+            graph_name=f"cluster_{initial_node.get_name()}_atomic",
+            label="",
+            peripheries=0,
+            cluster="true",
+        )
+        initial_subgraph.add_node(initial_node)
+        graph.add_subgraph(initial_subgraph)
+        graph.add_subgraph(atomic_states_subgraph)
+
+        initial = next(s for s in state.states if s.initial)
+        graph.add_edge(self._initial_edge(initial_node, initial))
+
+        for substate in state.states:
+            if substate.states:
+                subgraph = self._get_subgraph(substate)
+                self._graph_states(substate, subgraph)
+                graph.add_subgraph(subgraph)
+            else:
+                atomic_states_subgraph.add_node(self._state_as_node(substate))
+
+            for transition in substate.transitions:
                 if transition.internal:
                     continue
                 graph.add_edge(self._transition_as_edge(transition))
 
-        return graph
-
     def __call__(self):
         return self.get_graph()
 
@@ -165,7 +224,8 @@ def quickchart_write_svg(sm: StateMachine, path: str):
     >>> from tests.examples.order_control_machine import OrderControl
     >>> sm = OrderControl()
     >>> print(sm._graph().to_string())
-    digraph list {
+    digraph OrderControl {
+    compound=true;
     fontname=Arial;
     fontsize="10pt";
     label=OrderControl;

docs/images/order_control_machine_processing.png

@@ -32,6 +32,7 @@ def __init__(
         super().__init__(name, bases, attrs)
         registry.register(cls)
         cls.name = cls.__name__
+        cls.id = cls.name.lower()
         cls.states: States = States()
         cls.states_map: Dict[Any, State] = {}
         """Map of ``state.value`` to the corresponding :ref:`state`."""
@@ -46,6 +47,9 @@ def __init__(
 
         if not cls.states:
             return
+
+        cls._initials_by_document_order(cls.states)
+
         initials = [s for s in cls.states if s.initial]
         parallels = [s.id for s in cls.states if s.parallel]
         root_only_has_parallels = len(cls.states) == len(parallels)
@@ -59,11 +63,7 @@ def __init__(
             )
 
         try:
-            if root_only_has_parallels:
-                # TODO: Temp, whe should fix initial, and current state design
-                cls.initial_state: State = next(s for s in cls.states if s.initial)
-            else:
-                cls.initial_state: State = next(s for s in initials if s.initial)
+            cls.initial_state: State = next(s for s in initials if s.initial)
         except StopIteration:
             cls.initial_state = None  # Abstract SM still don't have states
 
@@ -77,6 +77,16 @@ def __init__(
 
         def __getattr__(self, attribute: str) -> Any: ...
 
+    def _initials_by_document_order(cls, states):
+        has_initial = False
+        for s in states:
+            cls._initials_by_document_order(s.states)
+            if s.initial:
+                has_initial = True
+                break
+        if not has_initial and states:
+            states[0]._initial = True
+
     def _check(cls):
         has_states = bool(cls.states)
         has_events = bool(cls._events)
@@ -233,7 +243,7 @@ def add_state(cls, id, state: State):
         for event in state.transitions.unique_events:
             cls.add_event(event)
 
-        for substate in state.substates:
+        for substate in state.states:
             cls.add_state(substate.id, substate)
 
     def add_event(cls, event, transitions=None):

docs/images/readme_trafficlightmachine.png

@@ -22,15 +22,15 @@ def __new__(  # type: ignore [misc]
         if not bases:
             return super().__new__(cls, classname, bases, attrs)  # type: ignore [return-value]
 
-        substates = []
+        states = []
         for key, value in attrs.items():
             if isinstance(value, State):
                 value._set_id(key)
-                substates.append(value)
+                states.append(value)
             if isinstance(value, TransitionList):
                 value.add_event(key)
 
-        return State(name=name, substates=substates, **kwargs)
+        return State(name=name, states=states, **kwargs)
 
 
 class State:
@@ -141,14 +141,15 @@ def __init__(
         initial: bool = False,
         final: bool = False,
         parallel: bool = False,
-        substates: Any = None,
+        states: Any = None,
         enter: Any = None,
         exit: Any = None,
     ):
         self.name = name
         self.value = value
         self.parallel = parallel
-        self.substates = substates or []
+        self.states = states or []
+        self.is_atomic = bool(not self.states)
         self._initial = initial
         self._final = final
         self._id: str = ""
@@ -161,12 +162,12 @@ def __init__(
         self.exit = self._specs.grouper(CallbackGroup.EXIT).add(
             exit, priority=CallbackPriority.INLINE
         )
-        self._init_substates()
+        self._init_states()
 
-    def _init_substates(self):
-        for substate in self.substates:
-            substate.parent = self
-            setattr(self, substate.id, substate)
+    def _init_states(self):
+        for state in self.states:
+            state.parent = self
+            setattr(self, state.id, state)
 
     def __eq__(self, other):
         return isinstance(other, State) and self.name == other.name and self.id == other.id
@@ -268,6 +269,7 @@ def __init__(
     ):
         self._state = ref(state)
         self._machine = ref(machine)
+        self._init_states()
 
     @property
     def name(self):
@@ -313,3 +315,15 @@ def id(self) -> str:
     @property
     def is_active(self):
         return self._machine().current_state == self
+
+    @property
+    def is_atomic(self):
+        return self._state().is_atomic
+
+    @property
+    def parent(self):
+        return self._state().parent
+
+    @property
+    def states(self):
+        return self._state().states

docs/images/test_state_machine_internal.png

@@ -61,6 +61,9 @@ def __getattr__(self, name: str):
     def __len__(self):
         return len(self._states)
 
+    def __getitem__(self, index):
+        return list(self)[index]
+
     def __iter__(self):
         return iter(self._states.values())