serverlessworkflow
diff --git a/‎serverlessworkflow/sdk/state_machine_extensions.py
Lines changed: 41 additions & 0 deletions b/‎serverlessworkflow/sdk/state_machine_extensions.py
Lines changed: 41 additions & 0 deletions
@@ -0,0 +1,41 @@
+from transitions.extensions.states import add_state_features, Tags, State
+from transitions.extensions import (
+    HierarchicalMachine,
+    GraphMachine,
+    HierarchicalGraphMachine,
+)
+
+
+class Metadata(State):
+    """Allows states to have metadata.
+    Attributes:
+        metadata (dict): A dictionary with the state metadata.
+    """
+
+    def __init__(self, *args, **kwargs):
+        """
+        Args:
+            **kwargs: If kwargs contains `metadata`, assign them to the attribute.
+        """
+        self.metadata = kwargs.pop("metadata", None)
+        super(Metadata, self).__init__(*args, **kwargs)
+
+    def __getattr__(self, key):
+        if value := self.metadata.get(key) is not None:
+            return value
+        return super(Metadata, self).__getattribute__(key)
+
+
+@add_state_features(Tags, Metadata)
+class CustomHierarchicalMachine(HierarchicalMachine):
+    pass
+
+
+@add_state_features(Tags, Metadata)
+class CustomHierarchicalGraphMachine(HierarchicalGraphMachine):
+    pass
+
+
+@add_state_features(Tags, Metadata)
+class CustomGraphMachine(GraphMachine):
+    pass