Separate out computation of power in PowerExprMixin into staticmethod

1 files changed, 32 insertions, 21 deletions

diff --git a/polymatrix/expression/mixins/powerexprmixin.py b/polymatrix/expression/mixins/powerexprmixin.py
index ffc9d07..4ca1b0b 100644
--- a/polymatrix/expression/mixins/powerexprmixin.py
+++ b/polymatrix/expression/mixins/powerexprmixin.py
@@ -28,40 +28,51 @@ class PowerExprMixin(ExpressionBaseMixin):
     @abstractmethod
     def right(self) -> ExpressionBaseMixin | int | float: ...
 
-    @override
-    def apply(self, state: ExpressionStateMixin) -> tuple[ExpressionStateMixin, PolyMatrixMixin]:
+    @staticmethod
+    def power(
+        state: ExpressionStateMixin,
+        left: ExpressionBaseMixin,
+        right: ExpressionBaseMixin | int | float
+    ) -> tuple[ExpressionStateMixin, PolyMatrixMixin]:
+        """ Compute the expression ```left ** right```. """
         exponent: int | None = None
 
         # Right (exponent) must end up being a scalar constant
-        if isinstance(self.right, int):
-            exponent = self.right
+        if isinstance(right, int):
+            exponent = right
             
-        elif isinstance(self.right, float):
-            exponent = int(self.right)
-            if not math.isclose(self.right - exponent, 0):
+        elif isinstance(right, float):
+            exponent = int(right)
+            if not math.isclose(right - exponent, 0):
                 raise ValueError("Cannot raise a variable to a non-integral power. "
-                                 f"Exponent {self.right} (float) is not close enough to an integer.")
+                                 f"Exponent {right} (float) is not close enough to an integer.")
 
-        elif isinstance(self.right, ExpressionBaseMixin):
-            state, right_polymatrix = self.right.apply(state)
-            right = right_polymatrix.at(0, 0).constant()
+        elif isinstance(right, ExpressionBaseMixin):
+            state, right_polymatrix = right.apply(state)
 
-            if not isinstance(right, int):
-                exponent = int(right)
-                if not math.isclose(right - exponent, 0):
+            r = right_polymatrix.at(0, 0).constant()
+            if not isinstance(r, int):
+                exponent = int(r)
+                if not math.isclose(r - exponent, 0):
                     raise ValueError("Cannot raise a variable to a non-integral power. "
-                                     f"Exponent {right}, resulting from {self.right} is not an integer.")
+                                     f"Exponent {right}, resulting from {right} is not an integer.")
             else:
-                exponent = right
+                exponent = r
 
         else:
-            raise TypeError(f"Cannot raise {self.left} to {self.right}, ",
-                            f"because exponet has type {type(self.right)}")
+            raise TypeError(f"Cannot raise {left} to {right}, ",
+                            f"because exponet has type {type(right)}")
 
-        state, left = self.left.apply(state)
-        result = left
+        state, base = left.apply(state)
+        result = base
         for _ in range(exponent -1):
-            state, result = ElemMultExprMixin.elem_mult(state, result, left)
+            state, result = ElemMultExprMixin.elem_mult(state, result, base)
 
         return state, result
 
+    @override
+    def apply(self, state: ExpressionStateMixin) -> tuple[ExpressionStateMixin, PolyMatrixMixin]:
+        # Separated into staticmethod like elem_mult for code resuse in
+        # CombinationsExprMixin
+        return self.power(state, self.left, self.right)
+