Make state.State a dataclass

Also, the members were (incorrectly) not instance variables

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}.")