Replace DATA_TYPE with FromDataTypes in expression module

5 files changed, 28 insertions, 23 deletions

diff --git a/polymatrix/expression/expression.py b/polymatrix/expression/expression.py
index e8f2e42..8d53354 100644
--- a/polymatrix/expression/expression.py
+++ b/polymatrix/expression/expression.py
@@ -11,7 +11,6 @@ from polymatrix.utils.getstacklines import get_stack_lines
 from polymatrix.polymatrix.abc import PolyMatrix
 from polymatrix.expressionstate.abc import ExpressionState
 from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
-from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
 from polymatrix.expression.op import (
     diff,
     integrate,
diff --git a/polymatrix/expression/from_.py b/polymatrix/expression/from_.py
index ac43069..011d5d7 100644
--- a/polymatrix/expression/from_.py
+++ b/polymatrix/expression/from_.py
@@ -1,5 +1,4 @@
-import numpy as np
-import sympy
+from polymatrix.expression.typing import FromDataTypes
 
 import polymatrix.expression.init
 
@@ -7,18 +6,6 @@ from polymatrix.expression.expression import init_expression, Expression
 from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
 from polymatrix.statemonad.abc import StateMonad
 
-# NP: how is this related to expression.init.DATA_TYPE?
-# FIXME: move to typing module
-FromDataTypes = (
-    str
-    | np.ndarray
-    | sympy.Matrix
-    | sympy.Expr
-    | tuple
-    | ExpressionBaseMixin
-    | StateMonad
-)
-
 
 # NP: this function name makes no sense to me, 
 # NP: also why does this allow None return type? init_expression always return
diff --git a/polymatrix/expression/init.py b/polymatrix/expression/init.py
index 75b63b6..0b4ced9 100644
--- a/polymatrix/expression/init.py
+++ b/polymatrix/expression/init.py
@@ -2,6 +2,8 @@ import typing
 import numpy as np
 import sympy
 
+from polymatrix.expression.typing import FromDataTypes
+
 import polymatrix.expression.impl
 
 from polymatrix.polymatrix.mixins import PolyMatrixMixin
@@ -154,14 +156,11 @@ def init_from_symmetric_matrix_expr(
     )
 
 
-# NP: replace with FromDataTypes?
-DATA_TYPE = str | np.ndarray | sympy.Matrix | sympy.Expr | tuple | ExpressionBaseMixin
-
 
 # NP: this function should be split up into smaller functions, one for each "from" type
 # NP: and each "from" should be documented, explaining how it is interpreted.
 def init_from_expr_or_none(
-    data: DATA_TYPE,
+    data: FromDataTypes,
 ) -> ExpressionBaseMixin | None:
     if isinstance(data, str):
         return init_parametrize_expr(
@@ -214,7 +213,7 @@ def init_from_expr_or_none(
     )
 
 
-def init_from_expr(data: DATA_TYPE):
+def init_from_expr(data: FromDataTypes):
     expr = init_from_expr_or_none(data)
 
     if expr is None:
diff --git a/polymatrix/expression/mixins/fromtupleexprmixin.py b/polymatrix/expression/mixins/fromtupleexprmixin.py
index e99742b..35c6e5d 100644
--- a/polymatrix/expression/mixins/fromtupleexprmixin.py
+++ b/polymatrix/expression/mixins/fromtupleexprmixin.py
@@ -3,6 +3,8 @@ import math
 import sympy
 import numpy as np
 
+from polymatrix.expression.typing import FromDataTypes
+
 from polymatrix.utils.getstacklines import FrameSummary
 from polymatrix.utils.tooperatorexception import to_operator_exception
 from polymatrix.polymatrix.init import init_poly_matrix
@@ -12,17 +14,16 @@ from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
 
 
 class FromTupleExprMixin(ExpressionBaseMixin):
-    DATA_TYPE = int | float | np.number | sympy.Expr | ExpressionBaseMixin
-
     @property
     @abc.abstractmethod
-    def data(self) -> tuple[tuple["FromTupleExprMixin.DATA_TYPE"]]: ...
+    def data(self) -> tuple[tuple[FromDataTypes]]: ...
 
     @property
     @abc.abstractmethod
     def stack(self) -> tuple[FrameSummary]: ...
 
     # overwrites the abstract method of `ExpressionBaseMixin`
+    # FIXME: break up this mess
     def apply(
         self,
         state: ExpressionStateMixin,
diff --git a/polymatrix/expression/typing.py b/polymatrix/expression/typing.py
new file mode 100644
index 0000000..58c5ca2
--- /dev/null
+++ b/polymatrix/expression/typing.py
@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
+from polymatrix.statemonad.abc import StateMonad
+
+import numpy.typing as npt
+import sympy
+
+
+FromDataTypes = (
+    str
+    | npt.NDArray
+    | sympy.Matrix
+    | sympy.Expr
+    | tuple
+    | ExpressionBaseMixin
+    | StateMonad
+)
+