User workflow machine

This machine binds the events to the User model, the StateChart is wrapped internally in the User class.

Demonstrates that multiple state machines can be used in the same model.

And that logic can be reused with listeners.

from dataclasses import dataclass
from enum import Enum

from statemachine.states import States

from statemachine import State
from statemachine import StateChart


class UserStatus(str, Enum):
    signup_incomplete = "SIGNUP_INCOMPLETE"
    signup_complete = "SIGNUP_COMPLETE"
    signup_rejected = "SIGNUP_REJECTED"
    operational_enabled = "OPERATIONAL_ENABLED"
    operational_disabled = "OPERATIONAL_DISABLED"
    operational_rescinded = "OPERATIONAL_RESCINDED"


class UserExperience(str, Enum):
    basic = "BASIC"
    premium = "PREMIUM"


@dataclass
class User:
    name: str
    email: str
    status: UserStatus = UserStatus.signup_incomplete
    experience: UserExperience = UserExperience.basic

    verified: bool = False

    def __post_init__(self):
        self._status_sm = UserStatusMachine(
            self, state_field="status", listeners=[MachineChangeListenter()]
        )
        self._status_sm.bind_events_to(self)

        self._experience_sm = UserExperienceMachine(
            self, state_field="experience", listeners=[MachineChangeListenter()]
        )
        self._experience_sm.bind_events_to(self)


class MachineChangeListenter:
    def before_transition(self, event: str, state: State):
        print(f"Before {event} in {state}")

    def on_enter_state(self, state: State, event: str):
        print(f"Entering {state} from {event}")


class UserStatusMachine(StateChart):
    catch_errors_as_events = False
    _states = States.from_enum(
        UserStatus,
        initial=UserStatus.signup_incomplete,
        final=[
            UserStatus.operational_rescinded,
            UserStatus.signup_rejected,
        ],
    )

    signup = _states.signup_incomplete.to(_states.signup_complete)
    reject = _states.signup_rejected.from_(
        _states.signup_incomplete,
        _states.signup_complete,
    )
    enable = _states.signup_complete.to(_states.operational_enabled)
    disable = _states.operational_enabled.to(_states.operational_disabled)
    rescind = _states.operational_rescinded.from_(
        _states.operational_enabled,
        _states.operational_disabled,
    )

    def on_signup(self, token: str):
        if token == "":
            raise ValueError("Token is required")
        self.model.verified = True  # type: ignore[union-attr]


class UserExperienceMachine(StateChart):
    _states = States.from_enum(
        UserExperience,
        initial=UserExperience.basic,
    )

    upgrade = _states.basic.to(_states.premium)
    downgrade = _states.premium.to(_states.basic)
  • user machine
  • user machine

Executing

def main():  # type: ignore[attr-defined]
    # By binding the events to the User model, the events can be fired directly from the model
    user = User(name="Frodo", email="frodo@lor.com")

    try:
        # Trying to signup with an empty token should raise an exception
        user.signup("")
    except Exception as e:
        print(e)

    assert user.verified is False

    user.signup("1234")

    assert user.status == UserStatus.signup_complete
    assert user.verified is True

    print(user.experience)
    user.upgrade()
    print(user.experience)


if __name__ == "__main__":
    main()
Before signup in Signup incomplete
Token is required
Before signup in Signup incomplete
Entering Signup complete from signup
UserExperience.basic
Before upgrade in Basic
Entering Premium from upgrade
UserExperience.premium