🚦 Finite State Machine (FSM)

Introduction

PBModular includes a built-in Finite State Machine (FSM) system that allows you to create handlers for multi-step user interactions, like registration forms, quizzes, or sequential command inputs. It works by tracking a "state" for each user in the main bot database.

Step 1: Defining a State Machine and States

1. Create a class that inherits from base.states.StateMachine.
2. Inside this class, define your states by creating class attributes that are instances of base.states.State.

Example:

# In main.py or a separate states.py file
from base.states import StateMachine, State

class RegistrationFSM(StateMachine):
    waiting_for_name = State()
    waiting_for_age = State()
    # The default state is None (no state set)

Step 2: Registering the State Machine

In your main BaseModule class, override the state_machine property to return your FSM class.

# In your BaseModule class (e.g., main.py)
from .my_fsm import RegistrationFSM # Assuming you defined it in my_fsm.py

class MyModule(BaseModule):
    # ... other properties ...

    @property
    def state_machine(self) -> type[StateMachine]:
        """Registers the FSM for this module."""
        return RegistrationFSM

Step 3: Creating State-Specific Handlers

Use the fsm_state parameter in your handler decorators (@command, @message, etc.) to specify which state(s) the handler should respond to.

• To handle a command that starts the process, you can omit fsm_state or set fsm_state=None, as users will be in the default (None) state.
• To handle subsequent messages, set fsm_state to the appropriate State object from your FSM class.

Step 4: Managing State in Handlers

When you register a StateMachine, your handlers can accept an additional fourth argument, typically named sm_controller. This object is an instance of your StateMachine class, bound to the current user, and provides methods to manage their state and data.

• sm_controller.get_state() -> Optional[str]: Get the user's current state name.
• sm_controller.set_state(State, data: Optional[dict] = None): Set the user's state. You can also store a dictionary of data along with it.
• sm_controller.get_data() -> dict: Get the user's stored data.
• sm_controller.update_data(**kwargs): Update the user's stored data dictionary.
• sm_controller.clear(keep_data: bool = False): Reset the user's state to default (None). By default, it also clears their data.

Complete Example:

# main.py
from base.module import BaseModule, command, message
from base.states import StateMachine, State
from pyrogram.types import Message

# 1. Define the FSM
class RegistrationFSM(StateMachine):
    waiting_for_name = State()
    waiting_for_age = State()

class FSMModule(BaseModule):
    # 2. Register the FSM
    @property
    def state_machine(self) -> type[StateMachine]:
        return RegistrationFSM

    # 3. Create handlers
    @command("register") # fsm_state=None by default
    async def start_registration(self, message: Message, sm_controller: RegistrationFSM):
        """Starts the user registration process."""
        await message.reply("Great! What is your name?")
        # 4. Set the next state
        await sm_controller.set_state(sm_controller.waiting_for_name)

    @message(fsm_state=RegistrationFSM.waiting_for_name)
    async def handle_name(self, message: Message, sm_controller: RegistrationFSM):
        """Handles the user's name input."""
        name = message.text
        await message.reply(f"Thanks, {name}! How old are you?")
        # Update data and set the next state
        await sm_controller.update_data(name=name)
        await sm_controller.set_state(sm_controller.waiting_for_age)

    @message(fsm_state=RegistrationFSM.waiting_for_age)
    async def handle_age(self, message: Message, sm_controller: RegistrationFSM):
        """Handles the user's age input."""
        if not message.text.isdigit():
            await message.reply("Please enter a valid number for your age.")
            return # Keep the user in the same state

        age = int(message.text)
        # Get all data and finalize
        user_data = await sm_controller.get_data()
        name = user_data.get("name")
        
        await message.reply(
            f"Registration complete!\n"
            f"Name: {name}\n"
            f"Age: {age}"
        )
        # Clear the state and data
        await sm_controller.clear()

    @command("cancel")
    async def cancel_process(self, message: Message, sm_controller: RegistrationFSM):
        """Cancels any ongoing FSM process."""
        current_state = await sm_controller.get_state()
        if current_state:
            await sm_controller.clear()
            await message.reply("Process cancelled.")
        else:
            await message.reply("Nothing to cancel.")