Rewrite state type annotations to avoid circular import

1 files changed, 12 insertions, 11 deletions

diff --git a/mdpoly/state.py b/mdpoly/state.py
index 56cd7e0..1ba44fe 100644
--- a/mdpoly/state.py
+++ b/mdpoly/state.py
@@ -1,18 +1,18 @@
-from .types import Number
-from .leaves import Var, Param
+from __future__ import annotations
+from typing import TYPE_CHECKING
 
-from typing import NewType
+if TYPE_CHECKING:
+    from .types import Number
+    from .leaves import Var, Param
 
-Index = NewType('Index', int)
 
-class State:
-    variables: dict[Var, Index]
-    parameters: dict[Param, Number]
+Index = int
+
 
-    def __init__(self, variables, parameters={}):
-        self._last_index = 0
-        self.variables = variables
-        self.parameters = parameters
+class State:
+    variables: dict[Var, Index] = {}
+    parameters: dict[Param, Number] = {}
+    _last_index: Index = -1 
 
     def _make_index(self) -> Index:
         """ Make a new index """
@@ -27,6 +27,7 @@ class State:
             return new_index
 
         return self.variables[var]
+