155 lines
6.4 KiB
Python
155 lines
6.4 KiB
Python
from enum import Enum
|
|
from typing import Optional, Callable, Protocol, Self, TypeVar, Generic, List, Union
|
|
|
|
|
|
# Anything hashable and equatable can be used as an input
|
|
class InputProtocol(Protocol):
|
|
def __eq__(self, __o: Self) -> bool:
|
|
...
|
|
|
|
def __hash__(self) -> int:
|
|
...
|
|
|
|
|
|
class TransitionInputEnum(Enum):
|
|
MATCH_REST = "#TRANSITION_MATCH_REST"
|
|
|
|
|
|
Input = TypeVar("Input", bound=InputProtocol)
|
|
TransitionInput = TypeVar("TransitionInput", bound=Union[InputProtocol, TransitionInputEnum])
|
|
|
|
|
|
class State:
|
|
def __init__(self, name: str) -> None:
|
|
self.__name = name
|
|
|
|
def get_name(self) -> str:
|
|
return self.__name
|
|
|
|
def __str__(self) -> str:
|
|
return f"[State {self.__name}]"
|
|
|
|
|
|
class Transition(Generic[TransitionInput]):
|
|
def __init__(self, transition_input: TransitionInput, origin_name: str, destination_name: str,
|
|
output_function: Optional[Callable[[Optional[object]], None]]) -> None:
|
|
self.__transition_input = transition_input
|
|
self.__origin_name = origin_name
|
|
self.__destination_name = destination_name
|
|
self.__destination_name_hash = hash(destination_name)
|
|
self.__output_function = output_function
|
|
|
|
def execute_output(self, context):
|
|
if self.__output_function is not None:
|
|
self.__output_function(context)
|
|
|
|
def get_destination_name_hash(self) -> int:
|
|
return self.__destination_name_hash
|
|
|
|
def __str__(self) -> str:
|
|
return f"[Transition ({self.__origin_name}, {self.__destination_name}, {self.__transition_input})]"
|
|
|
|
|
|
class Machine(Generic[Input]):
|
|
def __init__(self, initial_context: object = None) -> None:
|
|
if initial_context is None:
|
|
initial_context = {}
|
|
self.__start_state_hash: Optional[int] = None
|
|
self.__end_states: List[int] = []
|
|
self.__current_state_hash: Optional[int] = None
|
|
self.__states: dict[int, State] = {}
|
|
# dict[hash_origin, dict[hash_transition_input, Transition]]
|
|
self.__transitions: dict[int, dict[int, Transition]] = {}
|
|
self.__initial_context = initial_context
|
|
self.__context = initial_context
|
|
|
|
def add_transition(self, transition_input: TransitionInput, origin_name: str, destination_name: str,
|
|
output_function: Optional[Callable[[Optional[object]], None]]) -> None:
|
|
ho = hash(origin_name)
|
|
if ho not in self.__states.keys():
|
|
raise ValueError(f"Origin state {origin_name} does not exist.")
|
|
hd = hash(destination_name)
|
|
if hd not in self.__states.keys():
|
|
raise ValueError(f"Destination state {destination_name} does not exist.")
|
|
if isinstance(transition_input, TransitionInputEnum):
|
|
hi = hash(transition_input.name)
|
|
else:
|
|
hi = hash(transition_input)
|
|
t = Transition(transition_input, origin_name, destination_name, output_function)
|
|
if ho not in self.__transitions.keys():
|
|
self.__transitions[ho] = {hi: t}
|
|
elif hi not in self.__transitions[ho]:
|
|
self.__transitions[ho][hi] = t
|
|
elif hi in self.__transitions[ho]:
|
|
raise ValueError(f"Transition already exists for origin state {origin_name} and transition input {transition_input}.")
|
|
|
|
def add_state(self, state_name: str, is_start_state: bool = False, is_end_state: bool = False) -> None:
|
|
if state_name.startswith("#"):
|
|
raise ValueError("State names cannot start with # (reserved for special states)")
|
|
h: int = hash(state_name)
|
|
if is_start_state:
|
|
if self.__start_state_hash is None:
|
|
self.__start_state_hash = h
|
|
else:
|
|
raise ValueError("Adding a start state when one is already set.")
|
|
if is_end_state:
|
|
# Multiple end states possible
|
|
self.__end_states.append(h)
|
|
if h in self.__states.keys():
|
|
raise ValueError(f"State {state_name} already exists.")
|
|
state = State(state_name)
|
|
self.__states[h] = state
|
|
|
|
def process_input(self, i: Input) -> None:
|
|
if self.__start_state_hash is None:
|
|
raise ValueError("Processing input without a start state.")
|
|
if self.__current_state_hash is None:
|
|
self.__current_state_hash = self.__start_state_hash
|
|
# dict[hash_origin, dict[hash_transition_input, Transition]]
|
|
if self.__current_state_hash not in self.__transitions.keys():
|
|
raise ValueError("Processing input for a state with no transitions.")
|
|
assert self.__current_state_hash is not None
|
|
transitions = self.__transitions[self.__current_state_hash]
|
|
hi = hash(i)
|
|
transition: Optional[Transition] = None
|
|
if hi in transitions.keys():
|
|
# Transition found: execute and move to the destination state
|
|
transition = transitions[hi]
|
|
# No transition found, check if there's a special transition
|
|
if transition is None:
|
|
hc = hash(TransitionInputEnum.MATCH_REST)
|
|
if hc in transitions.keys():
|
|
transition = transitions[hc]
|
|
else:
|
|
raise ValueError(f"Found invalid transition while processing input {i} [{transition}].]")
|
|
assert transition is not None
|
|
transition.execute_output(self.__context)
|
|
self.__current_state_hash = transition.get_destination_name_hash()
|
|
|
|
def reset(self) -> None:
|
|
if self.__start_state_hash is None:
|
|
raise ValueError("Resetting a machine without starting state.")
|
|
self.__current_state_hash = self.__start_state_hash
|
|
self.__context = self.__initial_context
|
|
|
|
def get_current_state(self) -> Optional[State]:
|
|
if self.__current_state_hash is None:
|
|
return None # The machine has not started yet
|
|
return self.__states[self.__current_state_hash]
|
|
|
|
def is_in_final_state(self) -> bool:
|
|
if self.__current_state_hash is None:
|
|
return False
|
|
assert self.__current_state_hash is not None
|
|
return self.__current_state_hash in self.__end_states
|
|
|
|
def __str__(self) -> str:
|
|
r = [f"=MACHINE=", "\tSTATES"]
|
|
for state in self.__states.values():
|
|
r.append(f"\t\t{str(state)}")
|
|
r.append("\tTRANSITIONS")
|
|
for transition_origin in self.__transitions.keys():
|
|
for transition_input in self.__transitions[transition_origin]:
|
|
r.append(f"\t\t{str(self.__transitions[transition_origin][transition_input])}")
|
|
return "\n".join(r)
|