diff --git a/statemachine/statemachine.py b/statemachine/statemachine.py index 8d50ca77..943bd4ef 100644 --- 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 index e32a85a8..84395261 100644 --- 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):