Implementation of State Design Pattern in python

states.py

"""
Simple implementation of a finite state machine. 

Author: Tang U-Liang
Date: 2 Oct 2020 
"""
import attr 
from enum import Enum 
import inspect 


@attr.s
class StateBase:
    """ 
    Subclass StateBase to implement the transition and handling logic
    for actions. 

    To dispatch actions to the state, pass the handler name as a string,
    containing service, and other arguments to the dispatch method.
    """
    data = attr.ib(default=None)
    
    def dispatch(self, action, service, *args, **kwargs): 
        """
        The reason for returning the callable and it's bound arguments
        as seperate objects is so that we can defer execution of the effect
        until we have tested that the action can be handled. 

        This allows us to ensure the canonical sequence 
                exit => transition  => entry  
        of effects. 

        Under normal circumstances, a state action handler expects one to return a
        state object. To denote "internal transitions", explicitly annotate the return value 
        of the handler with None. This prevents the exit and entry actions from 
        firing and only the transition effect fires. 

        If the same state as current is returned, this is considered an external
        transition and the exit and entry actions *will* fire. 
        """
        effect = getattr(self, action, None)
        
        if effect: 
            sig = inspect.signature(effect)         
            ba = sig.bind(service, *args, **kwargs)
            return effect, sig, ba                   
        
        print(f"Unhandled action: {action}")
        return None, None, None

    def entry(self, service):
        """
        Override this method to implement an (optional) 
        entry effect
        """
        pass 

    def exit(self, service):
        """
        Override this method to implement an (optional) 
        exit effect
        """
        pass 


@attr.s
class ServiceBase: 
    """
    To implement a particular state machine, subclass ServiceBase
    and set an instance variable to desired state objects. 

    Services can then be initialized by passing name of initial
    state and an optional initial starting cache. 
    """
    
    _current:StateBase = attr.ib(init=False) 
    
    initial:str = attr.ib()
    cache = attr.ib(default=None) 

    def __attrs_post_init__(self): 
        self._current = getattr(self, self.initial)
    
    def send(self, action, *args, **kwargs): 

        effect, sig, ba = self._current.dispatch(action, self, *args, **kwargs)

        if effect:
            # Internal transitions 
            if sig.return_annotation is None: 
                effect(*ba.args, **ba.kwargs)

            # Normal (resp. external) state (resp. self) transitions 
            else: 
                self._current.exit(self)
                
                new_state = effect(*ba.args, **ba.kwargs)
                if new_state is None:
                    _msg = """
                        State transition handler is not annotated as None 
                        but failed to return a new state to transition to. 

                        Either return the current state or another state 
                        object. 
                        
                        If an internal transition is desired, annotate the 
                        return value with None. 
                        """
                    raise AssertionError(_msg)
                self._current = new_state
                self._current.entry(self)

traffic_light.py

"""
Simple implementation of a state machine that models a traffic light. 

Author: Tang U-Liang
Date: 2 Oct 2020 

>>> traffic_light.send('switch')
Leaving Red state
Entering Yellow state
>>> traffic_light.send('switch')
Leaving Yellow state
Entering Green state
>>> traffic_light.send('switch')
Leaving Green state
Entering Red state
>>> traffic_light.send('spoil')
Unhandled action: spoil
"""
from states import StateBase, ServiceBase 
import attr 


class Red(StateBase): 
    def switch(self, service):
        return service.YELLOW

    def exit(self, service):
        print("Leaving Red state")

    def entry(self, service): 
        print("Entering Red state")


class Yellow(StateBase):
    def switch(self, service):
        return service.GREEN

    def exit(self, service):
        print("Leaving Yellow state")

    def entry(self, service): 
        print("Entering Yellow state")



class Green(StateBase): 
    def switch(self, service): 
        return service.RED

    def exit(self, service):
        print("Leaving Green state")

    def entry(self, service): 
        print("Entering Green state")


@attr.s
class Service(ServiceBase): 
    RED = attr.ib(init=False, factory=Red)
    YELLOW = attr.ib(init=False, factory=Yellow)
    GREEN = attr.ib(init=False, factory=Green)
    

traffic_light = Service('RED')