Fixed typing of arithmetic methods

doc/release_notes.rst

@@ -12,6 +12,7 @@ Upcoming Version
   gap tolerance.
 * Improve the mapping of termination conditions for the SCIP solver
 * Treat GLPK's `integer undefined` status as not-OK
+* Fixed variable/expression arithmetic methods so that they correctly handle types
 
 Version 0.5.3
 --------------

linopy/expressions.py

@@ -12,10 +12,7 @@
 from collections.abc import Callable, Hashable, Iterator, Mapping, Sequence
 from dataclasses import dataclass, field
 from itertools import product, zip_longest
-from typing import (
-    TYPE_CHECKING,
-    Any,
-)
+from typing import TYPE_CHECKING, Any
 from warnings import warn
 
 import numpy as np
@@ -487,13 +484,16 @@
             )
             print(self)
 
-    def __add__(self, other: SideLike) -> LinearExpression:
+    def __add__(self, other: SideLike) -> LinearExpression | QuadraticExpression:
         """
         Add an expression to others.
 
         Note: If other is a numpy array or pandas object without axes names,
         dimension names of self will be filled in other
         """
+        if isinstance(other, QuadraticExpression):
+            return other.__add__(self)
+
         try:
             if np.isscalar(other):
                 return self.assign(const=self.const + other)
@@ -503,9 +503,8 @@
         except TypeError:
             return NotImplemented
 
-    def __radd__(self, other: int) -> LinearExpression | NotImplementedType:
-        # This is needed for using python's sum function
-        return self if other == 0 else NotImplemented
+    def __radd__(self, other: ConstantLike) -> LinearExpression:
+        return self.__add__(other)
 
     def __sub__(self, other: SideLike) -> LinearExpression:
         """
@@ -514,6 +513,9 @@
         Note: If other is a numpy array or pandas object without axes names,
         dimension names of self will be filled in other
         """
+        if isinstance(other, QuadraticExpression):
+            return other.__rsub__(self)
+
         try:
             if np.isscalar(other):
                 return self.assign_multiindex_safe(const=self.const - other)
@@ -523,7 +525,7 @@
         except TypeError:
             return NotImplemented
 
-    def __neg__(self) -> LinearExpression | QuadraticExpression:
+    def __neg__(self) -> LinearExpression:
         """
         Get the negative of the expression.
         """
@@ -536,14 +538,11 @@
         """
         Multiply the expr by a factor.
         """
+        if isinstance(other, QuadraticExpression):
+            return other.__rmul__(self)  # type: ignore
+
         try:
-            if isinstance(other, QuadraticExpression):
-                raise TypeError(
-                    "unsupported operand type(s) for *: "
-                    f"{type(self)} and {type(other)}. "
-                    "Higher order non-linear expressions are not yet supported."
-                )
-            elif isinstance(other, (variables.Variable, variables.ScalarVariable)):
+            if isinstance(other, (variables.Variable, variables.ScalarVariable)):
                 other = other.to_linexpr()
 
             if isinstance(other, (LinearExpression, ScalarLinearExpression)):
@@ -593,7 +592,7 @@
             raise ValueError("Power must be 2.")
         return self * self  # type: ignore
 
-    def __rmul__(self, other: ConstantLike) -> LinearExpression | QuadraticExpression:
+    def __rmul__(self, other: ConstantLike) -> LinearExpression:
         """
         Right-multiply the expr by a factor.
         """
@@ -1545,9 +1544,7 @@
         data = xr.Dataset(data.transpose(..., FACTOR_DIM, TERM_DIM))
         self._data = data
 
-    def __mul__(
-        self, other: ConstantLike | VariableLike | ExpressionLike
-    ) -> QuadraticExpression:
+    def __mul__(self, other: ConstantLike) -> QuadraticExpression:
         """
         Multiply the expr by a factor.
         """
@@ -1567,13 +1564,14 @@
             )
         return super().__mul__(other)  # type: ignore
 
+    def __rmul__(self, other: ConstantLike) -> QuadraticExpression:
+        return self.__mul__(other)
+
     @property
     def type(self) -> str:
         return "QuadraticExpression"
 
-    def __add__(
-        self, other: ConstantLike | VariableLike | ExpressionLike
-    ) -> QuadraticExpression:
+    def __add__(self, other: SideLike) -> QuadraticExpression:
         """
         Add an expression to others.
 
@@ -1592,21 +1590,13 @@
         except TypeError:
             return NotImplemented
 
-    def __radd__(
-        self, other: LinearExpression | int
-    ) -> LinearExpression | QuadraticExpression:
+    def __radd__(self, other: ConstantLike) -> QuadraticExpression:
         """
         Add others to expression.
         """
-        if type(other) is LinearExpression:
-            other = other.to_quadexpr()
-            return other.__add__(self)
-        elif other == 0:
-            return self
-        else:
-            return NotImplemented
+        return other.__add__(self)
 
-    def __sub__(self, other: SideLike | QuadraticExpression) -> QuadraticExpression:
+    def __sub__(self, other: SideLike) -> QuadraticExpression:
         """
         Subtract others from expression.
 
@@ -1624,15 +1614,17 @@
         except TypeError:
             return NotImplemented
 
-    def __rsub__(self, other: LinearExpression) -> QuadraticExpression:
+    def __rsub__(self, other: SideLike) -> QuadraticExpression:
         """
         Subtract expression from others.
         """
-        if type(other) is LinearExpression:
-            other = other.to_quadexpr()
-            return other.__sub__(self)
-        else:
-            return NotImplemented
+        return self.__neg__().__add__(other)
+
+    def __neg__(self) -> QuadraticExpression:
+        """
+        Get the negative of the expression.
+        """
+        return super().__neg__()  # type: ignore
 
     @property
     def solution(self) -> DataArray:
@@ -1875,7 +1867,7 @@
     A scalar linear expression container.
 
     In contrast to the LinearExpression class, a ScalarLinearExpression
-    only contains only one label. Use this class to create a constraint
+    only contains one label. Use this class to create a constraint
     in a rule.
     """
 

linopy/variables.py

@@ -382,22 +382,24 @@
         """
         return self.to_linexpr(-1)
 
-    def __mul__(
-        self, other: float | int | ndarray | Variable
-    ) -> LinearExpression | QuadraticExpression:
+    def __mul__(self, other: SideLike) -> LinearExpression | QuadraticExpression:
         """
-        Multiply variables with a coefficient.
+        Multiply variables with a coefficient, variable, or expression.
         """
         try:
-            if isinstance(
-                other, (expressions.LinearExpression, Variable, ScalarVariable)
-            ):
+            if isinstance(other, (Variable, ScalarVariable)):
                 return self.to_linexpr() * other
 
             return self.to_linexpr(other)
         except TypeError:
             return NotImplemented
 
+    def __rmul__(self, other: ConstantLike) -> LinearExpression:
+        """
+        Right-multiply variables by a constant
+        """
+        return self.to_linexpr(other)
+
     def __pow__(self, other: int) -> QuadraticExpression:
         """
         Power of the variables with a coefficient. The only coefficient allowed is 2.
@@ -407,15 +409,6 @@
             return expr._multiply_by_linear_expression(expr)
         return NotImplemented
 
-    def __rmul__(self, other: float | DataArray | int | ndarray) -> LinearExpression:
-        """
-        Right-multiply variables with a coefficient.
-        """
-        try:
-            return self.to_linexpr(other)
-        except TypeError:
-            return NotImplemented
-
     def __matmul__(
         self, other: LinearExpression | ndarray | Variable
     ) -> QuadraticExpression | LinearExpression:
@@ -449,9 +442,7 @@
         except TypeError:
             return NotImplemented
 
-    def __add__(
-        self, other: int | QuadraticExpression | LinearExpression | Variable
-    ) -> QuadraticExpression | LinearExpression:
+    def __add__(self, other: SideLike) -> LinearExpression:
         """
         Add variables to linear expressions or other variables.
         """
@@ -460,13 +451,13 @@
         except TypeError:
             return NotImplemented
 
-    def __radd__(self, other: int) -> Variable | NotImplementedType:
-        # This is needed for using python's sum function
-        return self if other == 0 else NotImplemented
+    def __radd__(self, other: ConstantLike) -> LinearExpression:
+        try:
+            return self.__add__(other)
+        except ValueError:
+            return NotImplemented
 
-    def __sub__(
-        self, other: QuadraticExpression | LinearExpression | Variable
-    ) -> QuadraticExpression | LinearExpression:
+    def __sub__(self, other: SideLike) -> LinearExpression:
         """
         Subtract linear expressions or other variables from the variables.
         """
@@ -475,6 +466,15 @@
         except TypeError:
             return NotImplemented
 
+    def __rsub__(self, other: ConstantLike) -> LinearExpression:
+        """
+        Subtract linear expressions or other variables from the variables.
+        """
+        try:
+            return self.to_linexpr(-1) + other
+        except TypeError:
+            return NotImplemented
+
     def __le__(self, other: SideLike) -> Constraint:
         return self.to_linexpr().__le__(other)
 

test/test_typing.py

@@ -0,0 +1,32 @@
+import xarray as xr
+from mypy import api
+
+import linopy
+
+
+def test_operations_with_data_arrays_are_typed_correctly() -> None:
+    m = linopy.Model()
+
+    a: xr.DataArray = xr.DataArray([1, 2, 3])
+
+    v: linopy.Variable = m.add_variables(lower=0.0, name="v")
+    e: linopy.LinearExpression = v * 1.0
+    q = v * v
+    assert isinstance(q, linopy.QuadraticExpression)
+
+    _ = a * v
+    _ = v * a
+    _ = v + a
+
+    _ = a * e
+    _ = e * a
+    _ = e + a
+
+    _ = a * q
+    _ = q * a
+    _ = q + a
+
+    # Get the path of this file
+    file_path = __file__
+    result = api.run([file_path])
+    assert result[2] == 0, "Mypy returned issues: " + result[0]