| 32 | |
| 33 | |
| 34 | class StateManager: |
| 35 | def __init__(self, state_cls: BaseState, init_args=None, init_kwargs=None): |
| 36 | self._state_cls = state_cls |
| 37 | self._init_args = init_args if init_args is not None else () |
| 38 | self._init_kwargs = init_kwargs if init_kwargs is not None else {} |
| 39 | self._state_cache = {} |
| 40 | self._current_context = None |
| 41 | |
| 42 | def get_state(self): |
| 43 | if self._current_context is None: |
| 44 | raise ValueError("No context is set. Please set a context before retrieving the state.") |
| 45 | if self._current_context not in self._state_cache.keys(): |
| 46 | self._state_cache[self._current_context] = self._state_cls(*self._init_args, **self._init_kwargs) |
| 47 | return self._state_cache[self._current_context] |
| 48 | |
| 49 | def set_context(self, name: str) -> None: |
| 50 | self._current_context = name |
| 51 | |
| 52 | def reset(self, *args, **kwargs) -> None: |
| 53 | for name, state in list(self._state_cache.items()): |
| 54 | state.reset(*args, **kwargs) |
| 55 | self._state_cache.pop(name) |
| 56 | self._current_context = None |
| 57 | |
| 58 | |
| 59 | class ModelHook: |
no outgoing calls
no test coverage detected
searching dependent graphs…