Remove expression dependency in expressionstate

1 files changed, 2 insertions, 8 deletions

diff --git a/polymatrix/expressionstate/mixins.py b/polymatrix/expressionstate/mixins.py
index 836032e..15f5712 100644
--- a/polymatrix/expressionstate/mixins.py
+++ b/polymatrix/expressionstate/mixins.py
@@ -6,7 +6,6 @@ from typing import NamedTuple, cast
 from math import prod
 import dataclasses
 
-from polymatrix.expression.expression import Expression
 from polymatrix.variable.abc import Variable
 
 from polymatrix.statemonad.mixins import StateCacheMixin
@@ -34,7 +33,7 @@ class ExpressionStateMixin(
         """ Map from variable objects to their indices. """
 
     def index(self, var: Variable) -> tuple[ExpressionStateMixin, IndexRange]:
-        """ Get the index of a variable. """
+        """ Get the index range of a variable. """
         # Unwrap if wrapped in expression object
         if not isinstance(var, Variable):
             raise ValueError("State can only index object of type Variable!")
@@ -47,13 +46,8 @@ class ExpressionStateMixin(
         size = prod(var.shape)
         index = IndexRange(start=self.n_variables, end=self.n_variables + size)
 
-        # FIXME: this is the only way the typechecker is happy, and it does
-        # make sense because ExpressionState is not a dataclass.
-        # The typechecker should be somehow informed that this is an ABC that
-        # will eventually become a dataclass, but of course it can't know that
-        from .impl import ExpressionStateImpl
         return dataclasses.replace(
-            cast(ExpressionStateImpl, self),
+            self,
             n_variables=self.n_variables + size,
             indices={**self.indices, var: index}
         ), index