Get initial state type from generic

This commit is contained in:
Thiago Moretto
2024-09-13 15:47:46 -03:00
parent ba8fbed30a
commit 2ef5cc37a9
2 changed files with 10 additions and 3 deletions

View File

@@ -115,6 +115,12 @@ class Flow(Generic[T], metaclass=FlowMeta):
_listeners: Dict[str, tuple[str, List[str]]] = {}
initial_state: Union[Type[T], T, None] = None
def __class_getitem__(cls, item):
print(f"[Flow.__class_getitem__] Getting initial state type: {item}")
class _FlowGeneric(cls):
_initial_state_T = item
return _FlowGeneric
def __init__(self):
print("[Flow.__init__] Initializing Flow")
self._methods: Dict[str, Callable] = {}
@@ -134,7 +140,9 @@ class Flow(Generic[T], metaclass=FlowMeta):
def _create_initial_state(self) -> T:
print("[Flow._create_initial_state] Creating initial state")
if self.initial_state is None:
if self.initial_state is None and hasattr(self, "_initial_state_T"):
return self._initial_state_T()
elif self.initial_state is None:
return {} # type: ignore
elif isinstance(self.initial_state, type):
return self.initial_state()

View File

@@ -9,8 +9,7 @@ class ExampleState(BaseModel):
message: str = ""
class StructuredExampleFlow(Flow):
initial_state = ExampleState
class StructuredExampleFlow(Flow[ExampleState]):
@start()
async def start_method(self):