diff options
| -rw-r--r-- | mdpoly/state.py | 13 |
1 files changed, 9 insertions, 4 deletions
diff --git a/mdpoly/state.py b/mdpoly/state.py index 1ba44fe..739ec7a 100644 --- a/mdpoly/state.py +++ b/mdpoly/state.py @@ -1,5 +1,6 @@ from __future__ import annotations from typing import TYPE_CHECKING +from dataclasses import dataclass, field if TYPE_CHECKING: from .types import Number @@ -9,9 +10,10 @@ if TYPE_CHECKING: Index = int +@dataclass class State: - variables: dict[Var, Index] = {} - parameters: dict[Param, Number] = {} + variables: dict[Var, Index] = field(default_factory=dict) + parameters: dict[Param, Number] = field(default_factory=dict) _last_index: Index = -1 def _make_index(self) -> Index: @@ -28,6 +30,9 @@ class State: return self.variables[var] - - + def from_index(self, index: Index) -> Var: + for var, idx in self.variables.items(): + if idx == index: + return var + raise IndexError(f"There is no variable with index {index}.") |