improve typing with PolynomialData and MonomialData

1 files changed, 7 insertions, 43 deletions

diff --git a/polymatrix/expression/utils/getvariableindices.py b/polymatrix/expression/utils/getvariableindices.py
index 4ee7037..d61da57 100644
--- a/polymatrix/expression/utils/getvariableindices.py
+++ b/polymatrix/expression/utils/getvariableindices.py
@@ -1,9 +1,14 @@
 import itertools
+import typing
 
+from polymatrix.expressionstate.abc import ExpressionState
 from polymatrix.expression.mixins.expressionbasemixin import ExpressionBaseMixin
 
 
-def get_variable_indices_from_variable(state, variable) -> tuple[int] | None:
+def get_variable_indices_from_variable(
+        state: ExpressionState, 
+        variable: ExpressionBaseMixin | int | typing.Any,
+) -> tuple[int, ...] | None:
     
     if isinstance(variable, ExpressionBaseMixin):
         state, variable_polynomial = variable.apply(state)
@@ -40,13 +45,12 @@ def get_variable_indices_from_variable(state, variable) -> tuple[int] | None:
     return state, variable_indices
 
 
+# not used, remove?
 def get_variable_indices(state, variables):
 
     if not isinstance(variables, tuple):
         variables = (variables,)
 
-    # assert isinstance(variables, tuple), f'{variables=}'
-
     def acc_variable_indices(acc, variable):
         state, indices = acc
 
@@ -61,43 +65,3 @@ def get_variable_indices(state, variables):
     )
 
     return state, indices
-
-
-    # global_state = [state]
-
-    # assert isinstance(variables, tuple), f'{variables=}'
-
-    # # if not isinstance(variables, tuple):
-    # #     variables = (variables,)
-
-    # def gen_indices():
-    #     for variable in variables:
-    #         if isinstance(variable, ExpressionBaseMixin):
-    #             global_state[0], variable_polynomial = variable.apply(global_state[0])
-
-    #             assert variable_polynomial.shape[1] == 1
-
-    #             for row in range(variable_polynomial.shape[0]):
-    #                 row_terms = variable_polynomial.get_poly(row, 0)
-
-    #                 assert len(row_terms) == 1, f'{row_terms} contains more than one term'
-                
-    #                 for monomial in row_terms.keys():
-    #                     assert len(monomial) <= 1, f'{monomial=} contains more than one variable'
-
-    #                     if len(monomial) == 0:
-    #                         continue
-                    
-    #                     assert monomial[0][1] == 1, f'{monomial[0]=}'
-    #                     yield monomial[0][0]
-
-    #         elif isinstance(variable, int):
-    #             yield variable
-
-    #         # else:
-    #         elif variable in global_state[0].offset_dict:
-    #             yield global_state[0].offset_dict[variable][0]
-
-    # indices = tuple(gen_indices())
-
-    # return global_state[0], indices