Search code examples
pythondecoratorfsm

Create class with decorator


I've implemented a Finite state machine in python. This works but implementing a state needs to write unnecessary code.

class State:
    def __init__(self):
        <do something>

    def __call__():
       <do something different>

class ConcreteState(State):
    def __init__(self):
       super().__init__()

    def __call__():
        super().__call__()
       <do concrete state implementation>

Is it possible to make a decorator in order to implement a concrete state like the following example?

@StateDecorator
def concreteState():
   <do concrete state implementation>

Solution

  • It's ugly, but since a decorator can return anything it's OK for it to return a class instead of a function:

    def StateDecorator(fn):
        class cls(State):
            def __call__(self):
                super().__call__()
                fn(self)
        return cls
    
    @StateDecorator
    def concreteState(self):
        print("cc", self)
    
    concreteState
    <class '__main__.cls'>
    

    Be aware that this could confuse any static analysis tools you're using.