diff options
-rw-r--r-- | polymatrix/expressionstate/mixins.py | 11 |
1 files changed, 8 insertions, 3 deletions
diff --git a/polymatrix/expressionstate/mixins.py b/polymatrix/expressionstate/mixins.py index ffaea0a..99b2cb1 100644 --- a/polymatrix/expressionstate/mixins.py +++ b/polymatrix/expressionstate/mixins.py @@ -2,7 +2,7 @@ from __future__ import annotations from typing import TYPE_CHECKING from abc import abstractmethod -from typing import NamedTuple +from typing import NamedTuple, cast from math import prod import dataclasses @@ -42,8 +42,13 @@ class ExpressionStateMixin( size = prod(var.shape) index = IndexRange(start=self.n_variables, end=self.n_variables + size) - return dataclasses.replace( # type: ignore[type-var] - self, + # FIXME: this is the only way the typechecker is happy, and it does + # make sense indeed 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), n_variables=self.n_variables + size, indices={**self.indices, var: index} ), index |