Simple pygame example with a basic state management

fsm_simple_example.py

from typing import TypeAlias, Type
import pygame
import sys

StateDict: TypeAlias = 'dict[str, StateBase | Type[StateBase]]'


class StateMachine:
    """Handles all the states and delegates to the correct one."""
    screen: pygame.Surface
    states: StateDict
    state: 'StateBase | None'

    def __init__(self, screen: pygame.Surface, states: StateDict, start_state: str):
        self.screen = screen
        self.states = states
        self.state = None
        self.change_state(start_state)

    def _try_call(self, method_name, *args, **kwargs):
        if self.state is None:
            raise Exception('No state set.')
        if hasattr(self.state, method_name):
            getattr(self.state, method_name)(*args, **kwargs)

    def handle_event(self, event):
        # def handle_event(self, event):
        self._try_call('handle_event', event)

    def update(self, delta: float):
        # def update(self, delta: float):
        self._try_call('update', delta)

    def draw(self):
        # def draw(self):
        self._try_call('draw')

    def change_state(self, state: str):
        if self.state is not None:
            # def on_exit(self, old_state: StateBase | None):
            self._try_call('on_exit', self.state)

        state_obj = self.states[state]
        if isinstance(state_obj, type):
            # if it's a class, changing to it will always create a new instance
            self.state = state_obj()
        else:
            self.state = state_obj

        self.state.fsm = self
        self.state.screen = self.screen

        # def on_enter(self):
        self._try_call('on_enter')


class StateBase:
    fsm: StateMachine
    screen: pygame.Surface


class SplashState(StateBase):
    """
    The first state that is shown when the game is started.
    A simple splash screen.
    """

    def __init__(self):
        self.start_time = None
        self.show_time = 2

        self.font = pygame.font.Font('Quicksand/static/Quicksand-Regular.ttf', 100)
        self.text = self.font.render('Splash!', True, (255, 255, 255))

    def handle_event(self, event):
        pass

    def update(self, delta: float):
        if self.start_time is None:
            self.start_time = pygame.time.get_ticks()

        if pygame.time.get_ticks() - self.start_time > self.show_time * 1000:
            self.fsm.change_state('menu')

    def draw(self):
        text_rect = self.text.get_rect()
        text_rect.center = self.screen.get_rect().center
        self.screen.blit(self.text, text_rect)


class MenuState(StateBase):
    """
    The main menu of the game.
    """

    def __init__(self):
        self.rotation = 0

        self.box = pygame.Surface((100, 100), pygame.SRCALPHA)
        self.box.fill((255, 0, 0))

    def handle_event(self, event):
        # set rotation based on mouse x position
        if event.type == pygame.MOUSEMOTION:
            self.rotation = event.pos[0] / 1360 * 360 * -1

    def update(self, delta: float):
        pass

    def draw(self):
        # rotate the box
        rotated_box = pygame.transform.rotate(self.box, self.rotation)
        rotated_box_rect = rotated_box.get_rect()
        rotated_box_rect.center = self.screen.get_rect().center

        self.screen.blit(rotated_box, rotated_box_rect)


def main():
    pygame.init()
    screen = pygame.display.set_mode((1360, 705))
    clock = pygame.time.Clock()

    fsm = StateMachine(screen, {
        'splash': SplashState,
        'menu': MenuState()
    }, 'splash')

    while fsm.state is not None:
        delta = clock.tick(60) / 1000
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()

            fsm.handle_event(event)

        screen.fill((0, 0, 0))
        fsm.update(delta)
        fsm.draw()
        pygame.display.update()


if __name__ == '__main__':
    main()