How to add tags to an existing state object in Pytransitions?

Question

In the project I'm assigned to we use pytransitions. Our states are created, get equipped with additional attributes and added to a list one by one as objects first. Then this list of State objects is passed to a Machine object. Here's a simple example:

from transitions import State

states = []
state_initial = State("initial", on_exit="some_callback")
text = "this is some text"
state.text = text
states.append(state)

This is how a machine is created:

from transitions import Machine
from some_library import SomeClass
from my_module import user_transitions

class User:
    states = states
    initial_state = states[0]
    def __init__(self, some_param: str, another_param: SomeClass = default_param):
        self.machine = Machine(model=self,
                               states=User.states,
                               initial=User.initial_state,
                               transitions=user_transitions,
                               prepare_event="preparing_callback",
                               after_state_change="ending_callback")

What I'd like to do is to add tags to my states at the time of or after state object creation. I mean the tags in transitions.extensions.states, so I could get them with is_tag kind of methods like in docs. Something like state_initial.add_tags(["tag1", "tag2"]) or state_initial = State("initial", on_exit="some_callback", tags=["tag1", "tag2"]) or in any other way considering my legacy setup. How do I go about this?

Answer 1

My first suggestion would be to check whether you can streamline the state creation process by using a dedicated TextState instead of just assigning an additional attribute. This way you can keep your state configuration a bit more comprehensible. Reading machine configurations from yaml or json files gets way easier as well.

from transitions import Machine, State
from transitions.extensions.states import Tags

# option a) create a custom state class and use it by default
# class TextState and CustomState could be combined of course
# splitting CustomState into two classes decouples tags from the 
# original state creation code
class TextState(State):

    def __init__(self, *args, **kwargs):
        self.text = kwargs.pop('text', '')
        super(TextState, self).__init__(*args, **kwargs)

class CustomState(Tags, TextState):
    pass


class CustomMachine(Machine):
    state_cls = CustomState


states = []
state_initial = CustomState("initial", text="this is some text")
# we can pass tags for initialization
state_foo = dict(name="foo", text="bar!", tags=['success'])
states.append(state_initial)
states.append(state_foo)

# [...] CustomMachine(model=self, states=User.states, initial=User.initial_state)

But your question was about how you can inject tag capability AFTER states have been created. Probably because it would need major refactoring and deep digging to alter state creation. Adding state.tags = ['your', 'tags', 'here'] is fine and should work out of the box for graph and markup creation. To get state.is_<tag> working you can alter its __class__ attribute:

from transitions import Machine, State
from transitions.extensions.states import Tags

# option b) patch __class__
states = []
state_initial = State("initial")
state_initial.text = "this is some text"
# we can pass tags for initialization
state_foo = State("foo")
state_foo.text = "bar!"
state_foo.tags = ['success']
states.append(state_initial)
states.append(state_foo)

# patch all states
for s in states:
    s.__class__ = Tags
    s.tags = []

# add tag to state_foo
states[1].tags.append('success')

class User:

    states = states
    initial_state = states[0]

    def __init__(self):
        self.machine = Machine(model=self,
                               states=User.states,
                               initial=User.initial_state)

user = User()
user.to_foo()
assert user.machine.get_state(user.state).is_success  # works!
assert not user.machine.get_state(user.state).is_superhero  # bummer...

But again, from my experience code becomes much more comprehensible and reusable when you strive to separate machine configuration from the rest of the code base. Patching states somewhere in the code and assigning custom paramters might be overlooked by the next guy working with your code and it surely is surprising when states change their class between two debugging breakpoints.