chore: Improve observers visitor reusing the initial setup loop

fgmacedo · Oct 6, 2023 · a17fe7f · a17fe7f
1 parent e209c03
commit a17fe7f
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 18 deletions.
diff --git a/statemachine/statemachine.py b/statemachine/statemachine.py
@@ -154,10 +154,32 @@ def _setup(self, initial_transition: Transition):
         default_resolver = resolver_factory(machine, model)
 
         register = partial(self._callbacks_registry.register, resolver=default_resolver)
-        self._visit_states_and_transitions(lambda x: x._setup(register))
+
+        observer_visitor = self._build_observers_visitor(machine, model)
+
+        def setup_visitor(visited):
+            visited._setup(register)
+            observer_visitor(visited)
+
+        self._visit_states_and_transitions(setup_visitor)
 
         initial_transition._setup(register)
-        self.add_observer(machine, model)
+
+    def _build_observers_visitor(self, *observers):
+        registry_callbacks = [
+            (
+                self._callbacks_registry.build_register_function_for_resolver(
+                    resolver_factory(observer)
+                )
+            )
+            for observer in observers
+        ]
+
+        def _add_observer_for_resolver(visited):
+            for register in registry_callbacks:
+                visited._add_observer(register)
+
+        return _add_observer_for_resolver
 
     def add_observer(self, *observers):
         """Add an observer.
@@ -170,18 +192,8 @@ def add_observer(self, *observers):
             :ref:`observers`.
         """
 
-        resolvers = [resolver_factory(ObjectConfig.from_obj(o)) for o in observers]
-
-        def _add_observer_for_resolver(x):
-            for resolver in resolvers:
-                register = (
-                    self._callbacks_registry.build_register_function_for_resolver(
-                        resolver
-                    )
-                )
-                x._add_observer(register)
-
-        self._visit_states_and_transitions(lambda x: _add_observer_for_resolver(x))
+        visitor = self._build_observers_visitor(*observers)
+        self._visit_states_and_transitions(visitor)
         return self
 
     def _repr_html_(self):

diff --git a/tests/test_profiling.py b/tests/test_profiling.py
@@ -43,10 +43,6 @@ def after_receive_payment(self):
 def exercise_order():
     order = Order()
     order.state_machine.waiting_for_payment.is_active
-    # order.state_machine.add_to_order(10)
-    # order.state_machine.receive_payment(10)
-    # order.state_machine.process_order()
-    # order.state_machine.ship_order()
 
 
 def test_setup_performance(benchmark):