Changes to typing

Author: Robbie Muir

linopy/expressions.py

@@ -12,7 +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, TypeVar, overload
 from warnings import warn
 
 import numpy as np
@@ -531,10 +531,18 @@ def __neg__(self) -> LinearExpression:
         """
         return self.assign_multiindex_safe(coeffs=-self.coeffs, const=-self.const)
 
+    @overload
     def __mul__(
-        self,
+        self: GenericLinearExpression, other: ConstantLike
+    ) -> GenericLinearExpression: ...
+
+    @overload
+    def __mul__(self, other: VariableLike | ExpressionLike) -> QuadraticExpression: ...
+
+    def __mul__(
+        self: GenericLinearExpression,
         other: SideLike,
-    ) -> LinearExpression | QuadraticExpression:
+    ) -> GenericLinearExpression | QuadraticExpression:
         """
         Multiply the expr by a factor.
         """
@@ -577,7 +585,9 @@ def _multiply_by_linear_expression(
             res = res + self.reset_const() * other.const
         return res  # type: ignore
 
-    def _multiply_by_constant(self, other: ConstantLike) -> LinearExpression:
+    def _multiply_by_constant(
+        self: GenericLinearExpression, other: ConstantLike
+    ) -> GenericLinearExpression:
         multiplier = as_dataarray(other, coords=self.coords, dims=self.coord_dims)
         coeffs = self.coeffs * multiplier
         assert all(coeffs.sizes[d] == s for d, s in self.coeffs.sizes.items())
@@ -592,7 +602,9 @@ def __pow__(self, other: int) -> QuadraticExpression:
             raise ValueError("Power must be 2.")
         return self * self  # type: ignore
 
-    def __rmul__(self, other: ConstantLike) -> LinearExpression:
+    def __rmul__(
+        self: GenericLinearExpression, other: ConstantLike
+    ) -> GenericLinearExpression:
         """
         Right-multiply the expr by a factor.
         """
@@ -611,11 +623,18 @@ def __matmul__(
         return (self * other).sum(dim=common_dims)
 
     def __div__(
-        self, other: Variable | ConstantLike
-    ) -> LinearExpression | QuadraticExpression:
+        self: GenericLinearExpression, other: SideLike
+    ) -> GenericLinearExpression:
         try:
             if isinstance(
-                other, (LinearExpression, variables.Variable, variables.ScalarVariable)
+                other,
+                (
+                    variables.Variable,
+                    variables.ScalarVariable,
+                    LinearExpression,
+                    ScalarLinearExpression,
+                    QuadraticExpression,
+                ),
             ):
                 raise TypeError(
                     "unsupported operand type(s) for /: "
@@ -627,8 +646,8 @@ def __div__(
             return NotImplemented
 
     def __truediv__(
-        self, other: Variable | ConstantLike
-    ) -> LinearExpression | QuadraticExpression:
+        self: GenericLinearExpression, other: SideLike
+    ) -> GenericLinearExpression:
         return self.__div__(other)
 
     def __le__(self, rhs: SideLike) -> Constraint:
@@ -1514,6 +1533,9 @@ def to_polars(self) -> pl.DataFrame:
     iterate_slices = iterate_slices
 
 
+GenericLinearExpression = TypeVar("GenericLinearExpression", bound=LinearExpression)
+
+
 class QuadraticExpression(LinearExpression):
     """
     A quadratic expression consisting of terms of coefficients and variables.
@@ -1544,7 +1566,7 @@ def __init__(self, data: Dataset | None, model: Model) -> None:
         data = xr.Dataset(data.transpose(..., FACTOR_DIM, TERM_DIM))
         self._data = data
 
-    def __mul__(self, other: ConstantLike) -> QuadraticExpression:
+    def __mul__(self, other: SideLike) -> QuadraticExpression:
         """
         Multiply the expr by a factor.
         """
@@ -1553,6 +1575,7 @@ def __mul__(self, other: ConstantLike) -> QuadraticExpression:
             (
                 LinearExpression,
                 QuadraticExpression,
+                ScalarLinearExpression,
                 variables.Variable,
                 variables.ScalarVariable,
             ),
@@ -1562,9 +1585,9 @@ def __mul__(self, other: ConstantLike) -> QuadraticExpression:
                 f"{type(self)} and {type(other)}. "
                 "Higher order non-linear expressions are not yet supported."
             )
-        return super().__mul__(other)  # type: ignore
+        return super().__mul__(other)
 
-    def __rmul__(self, other: ConstantLike) -> QuadraticExpression:
+    def __rmul__(self, other: SideLike) -> QuadraticExpression:
         return self.__mul__(other)
 
     @property

linopy/variables.py

@@ -52,7 +52,14 @@
 )
 from linopy.config import options
 from linopy.constants import HELPER_DIMS, TERM_DIM
-from linopy.types import ConstantLike, DimsLike, NotImplementedType, SideLike
+from linopy.types import (
+    ConstantLike,
+    DimsLike,
+    ExpressionLike,
+    NotImplementedType,
+    SideLike,
+    VariableLike,
+)
 
 if TYPE_CHECKING:
     from linopy.constraints import AnonymousScalarConstraint, Constraint
@@ -382,7 +389,13 @@ def __neg__(self) -> LinearExpression:
         """
         return self.to_linexpr(-1)
 
-    def __mul__(self, other: SideLike) -> LinearExpression | QuadraticExpression:
+    @overload
+    def __mul__(self, other: ConstantLike) -> LinearExpression: ...
+
+    @overload
+    def __mul__(self, other: ExpressionLike | VariableLike) -> QuadraticExpression: ...
+
+    def __mul__(self, other: SideLike) -> ExpressionLike:
         """
         Multiply variables with a coefficient, variable, or expression.
         """
@@ -398,7 +411,7 @@ def __rmul__(self, other: ConstantLike) -> LinearExpression:
         """
         Right-multiply variables by a constant
         """
-        return self.to_linexpr(other)
+        return self * other
 
     def __pow__(self, other: int) -> QuadraticExpression:
         """
@@ -1539,6 +1552,8 @@ def __mul__(self, coeff: int | float) -> ScalarLinearExpression:
         return self.to_scalar_linexpr(coeff)
 
     def __rmul__(self, coeff: int | float) -> ScalarLinearExpression:
+        if isinstance(coeff, Variable):
+            return NotImplemented
         return self.to_scalar_linexpr(coeff)
 
     def __div__(self, coeff: int | float) -> ScalarLinearExpression:

test/test_compatible_arithmetrics.py

@@ -97,7 +97,7 @@ def test_arithmetric_operations_variable(m: Model) -> None:
     assert_linequal(x + data, x + other_datatype)
     assert_linequal(x - data, x - other_datatype)
     assert_linequal(x * data, x * other_datatype)
-    assert_linequal(x / data, x / other_datatype)
+    assert_linequal(x / data, x / other_datatype)  # type: ignore
     assert_linequal(data * x, other_datatype * x)  # type: ignore
     assert x.__add__(object()) is NotImplemented
     assert x.__sub__(object()) is NotImplemented
@@ -131,7 +131,7 @@ def test_arithmetric_operations_vars_and_expr(m: Model) -> None:
     assert_linequal(x**2 + x, x + x**2)
     assert_linequal(x**2 * 2, x**2 * 2)
     with pytest.raises(TypeError):
-        _ = x**2 * x  # type: ignore
+        _ = x**2 * x
 
 
 def test_arithmetric_operations_con(m: Model) -> None:

test/test_optimization.py

@@ -386,7 +386,7 @@ def test_default_setting_expression_sol_accessor(
     qexpr = 4 * x**2
     assert_equal(qexpr.solution, 4 * x.solution**2)
 
-    qexpr = 4 * x * y
+    qexpr = 4 * (x * y)  # type: ignore
     assert_equal(qexpr.solution, 4 * x.solution * y.solution)
 
 

test/test_typing.py

@@ -5,14 +5,18 @@
 
 
 def test_operations_with_data_arrays_are_typed_correctly() -> None:
+    # Get the path of this file
+    file_path = __file__
+    result = api.run([file_path])
+    assert result[2] == 0, "Mypy returned issues: " + result[0]
+
     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
@@ -25,8 +29,3 @@ def test_operations_with_data_arrays_are_typed_correctly() -> None:
     _ = 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]