Further typing changes

Author: Robbie Muir

.gitignore

@@ -34,3 +34,6 @@ benchmark/notebooks/.ipynb_checkpoints
 benchmark/scripts/__pycache__
 benchmark/scripts/benchmarks-pypsa-eur/__pycache__
 benchmark/scripts/leftovers/
+
+# IDE
+.idea/

linopy/expressions.py

@@ -484,6 +484,14 @@ def print(self, display_max_rows: int = 20, display_max_terms: int = 20) -> None
             )
             print(self)
 
+    @overload
+    def __add__(
+        self, other: ConstantLike | Variable | ScalarLinearExpression | LinearExpression
+    ) -> LinearExpression: ...
+
+    @overload
+    def __add__(self, other: QuadraticExpression) -> QuadraticExpression: ...
+
     def __add__(self, other: SideLike) -> LinearExpression | QuadraticExpression:
         """
         Add an expression to others.
@@ -506,24 +514,16 @@ def __add__(self, other: SideLike) -> LinearExpression | QuadraticExpression:
     def __radd__(self, other: ConstantLike) -> LinearExpression:
         return self.__add__(other)
 
-    def __sub__(self, other: SideLike) -> LinearExpression:
-        """
-        Subtract others from expression.
-
-        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)
+    @overload
+    def __sub__(
+        self, other: ConstantLike | Variable | ScalarLinearExpression | LinearExpression
+    ) -> LinearExpression: ...
 
-        try:
-            if np.isscalar(other):
-                return self.assign_multiindex_safe(const=self.const - other)
+    @overload
+    def __sub__(self, other: QuadraticExpression) -> QuadraticExpression: ...
 
-            other = as_expression(other, model=self.model, dims=self.coord_dims)
-            return merge([self, -other], cls=self.__class__)
-        except TypeError:
-            return NotImplemented
+    def __sub__(self, other: SideLike) -> LinearExpression | QuadraticExpression:
+        return self.__add__(-other)
 
     def __neg__(self) -> LinearExpression:
         """

linopy/variables.py

@@ -455,7 +455,15 @@ def __truediv__(
         except TypeError:
             return NotImplemented
 
-    def __add__(self, other: SideLike) -> LinearExpression:
+    @overload
+    def __add__(
+        self, other: ConstantLike | Variable | ScalarLinearExpression | LinearExpression
+    ) -> LinearExpression: ...
+
+    @overload
+    def __add__(self, other: QuadraticExpression) -> QuadraticExpression: ...
+
+    def __add__(self, other: SideLike) -> LinearExpression | QuadraticExpression:
         """
         Add variables to linear expressions or other variables.
         """
@@ -470,7 +478,15 @@ def __radd__(self, other: ConstantLike) -> LinearExpression:
         except ValueError:
             return NotImplemented
 
-    def __sub__(self, other: SideLike) -> LinearExpression:
+    @overload
+    def __sub__(
+        self, other: ConstantLike | Variable | ScalarLinearExpression | LinearExpression
+    ) -> LinearExpression: ...
+
+    @overload
+    def __sub__(self, other: QuadraticExpression) -> QuadraticExpression: ...
+
+    def __sub__(self, other: SideLike) -> LinearExpression | QuadraticExpression:
         """
         Subtract linear expressions or other variables from the variables.
         """

test/test_linear_expression.py

@@ -226,6 +226,14 @@ def test_linear_expression_with_addition(m: Model, x: Variable, y: Variable) ->
     assert_linequal(expr, expr2)
 
 
+def test_linear_expression_with_raddition(m: Model, x: Variable):
+    expr = x * 1.0
+    expr_2: LinearExpression = 10.0 + expr  # type: ignore
+    assert isinstance(expr, LinearExpression)
+    expr_3: LinearExpression = expr + 10.0  # type: ignore
+    assert_linequal(expr_2, expr_3)
+
+
 def test_linear_expression_with_subtraction(m: Model, x: Variable, y: Variable) -> None:
     expr = x - y
     assert isinstance(expr, LinearExpression)

test/test_quadratic_expression.py

@@ -144,8 +144,12 @@ def test_quadratic_expression_raddition(x: Variable, y: Variable) -> None:
     assert (expr.const == 5).all()
     assert expr.nterm == 2
 
-    with pytest.raises(TypeError):
-        5 + x * y + x
+    expr_2 = 5 + x * y + x
+    assert isinstance(expr_2, QuadraticExpression)
+    assert (expr_2.const == 5).all()
+    assert expr_2.nterm == 2
+
+    assert_quadequal(expr, expr_2)
 
 
 def test_quadratic_expression_subtraction(x: Variable, y: Variable) -> None: