Annotate arithmetic operations

.github/workflows/stubs.yml

@@ -90,7 +90,6 @@ jobs:
       - name: Install Ruff
         uses: astral-sh/ruff-action@v3
         with:
-          version: "0.14.11"
           args: "--version"
 
       - name: Lint type stubs

src/pyscipopt/scip.pyi

@@ -4,7 +4,7 @@ from typing import Any, AnyStr, ClassVar, Literal, SupportsFloat, Union, overloa
 
 import numpy as np
 from _typeshed import Incomplete
-from typing_extensions import CapsuleType, disjoint_base
+from typing_extensions import CapsuleType, Self, disjoint_base
 
 CONST: Term
 EventNames: dict
@@ -20,21 +20,37 @@ _core_dot: Incomplete
 _core_dot_2d: Incomplete
 _core_sum: Incomplete
 buildGenExprObj: Incomplete
-cos: Incomplete
-exp: Incomplete
 expr_to_array: Incomplete
 expr_to_nodes: Incomplete
 is_memory_freed: Incomplete
-log: Incomplete
 print_memory_in_use: Incomplete
 quickprod: Incomplete
 quicksum: Incomplete
 readStatistics: Incomplete
-sin: Incomplete
-sqrt: Incomplete
 str_conversion: Incomplete
 value_to_array: Incomplete
 
+@overload
+def exp(x: float | Expr | GenExpr) -> UnaryExpr: ...
+@overload
+def exp(x: list | tuple | np.ndarray | MatrixExpr) -> MatrixGenExpr: ...
+@overload
+def log(x: float | Expr | GenExpr) -> UnaryExpr: ...
+@overload
+def log(x: list | tuple | np.ndarray | MatrixExpr) -> MatrixGenExpr: ...
+@overload
+def sqrt(x: float | Expr | GenExpr) -> UnaryExpr: ...
+@overload
+def sqrt(x: list | tuple | np.ndarray | MatrixExpr) -> MatrixGenExpr: ...
+@overload
+def sin(x: float | Expr | GenExpr) -> UnaryExpr: ...
+@overload
+def sin(x: list | tuple | np.ndarray | MatrixExpr) -> MatrixGenExpr: ...
+@overload
+def cos(x: float | Expr | GenExpr) -> UnaryExpr: ...
+@overload
+def cos(x: list | tuple | np.ndarray | MatrixExpr) -> MatrixGenExpr: ...
+
 @disjoint_base
 class Benders:
     model: Incomplete
@@ -236,6 +252,9 @@ class Conshdlr:
 class Constant(GenExpr):
     number: Incomplete
     def __init__(self, *args: Incomplete, **kwargs: Incomplete) -> None: ...
+    def __pow__(  # type: ignore[override]
+        self, other: float | Constant, mod: Incomplete = ..., /
+    ) -> Constant: ...
 
 @disjoint_base
 class Constraint:
@@ -335,40 +354,75 @@ class ExprLike:
         *args: Incomplete,
         **kwargs: Incomplete,
     ) -> Incomplete: ...
-    def __radd__(self, other: object, /) -> Incomplete: ...
-    def __sub__(self, other: object, /) -> Incomplete: ...
-    def __rsub__(self, other: object, /) -> Incomplete: ...
-    def __rmul__(self, other: object, /) -> Incomplete: ...
-    def __rtruediv__(self, other: object, /) -> GenExpr: ...
-    def __neg__(self, /) -> Union[Expr, GenExpr]: ...
-    def __pos__(self, /) -> Union[Expr, GenExpr]: ...
-    def __abs__(self) -> GenExpr: ...
-    def exp(self) -> GenExpr: ...
-    def log(self) -> GenExpr: ...
-    def sqrt(self) -> GenExpr: ...
-    def sin(self) -> GenExpr: ...
-    def cos(self) -> GenExpr: ...
+    def __pos__(self, /) -> Self: ...
+    def __abs__(self, /) -> UnaryExpr: ...
+    def exp(self) -> UnaryExpr: ...
+    def log(self) -> UnaryExpr: ...
+    def sqrt(self) -> UnaryExpr: ...
+    def sin(self) -> UnaryExpr: ...
+    def cos(self) -> UnaryExpr: ...
+    @overload
+    def __le__(self, other: float | ExprLike, /) -> ExprCons: ...
+    @overload
+    def __le__(self, other: np.ndarray | MatrixExpr, /) -> MatrixExprCons: ...  # type: ignore[misc]
+    @overload
+    def __ge__(self, other: float | ExprLike, /) -> ExprCons: ...
+    @overload
+    def __ge__(self, other: np.ndarray | MatrixExpr, /) -> MatrixExprCons: ...  # type: ignore[misc]
+    @overload  # type: ignore[override]
+    def __eq__(self, other: float | ExprLike, /) -> ExprCons: ...
+    @overload
+    def __eq__(self, other: np.ndarray | MatrixExpr, /) -> MatrixExprCons: ...
 
 @disjoint_base
 class Expr(ExprLike):
     terms: Incomplete
     def __init__(self, terms: Incomplete = ...) -> None: ...
     def degree(self) -> Incomplete: ...
     def normalize(self) -> Incomplete: ...
-    def __add__(self, other: Incomplete, /) -> Incomplete: ...
-    def __eq__(self, other: object, /) -> bool: ...
-    def __ge__(self, other: object, /) -> bool: ...
+    def __neg__(self, /) -> Expr: ...
+    @overload
+    def __add__(self, other: float | Expr, /) -> Expr: ...
+    @overload
+    def __add__(self, other: GenExpr, /) -> SumExpr: ...
+    @overload
+    def __add__(self, other: np.ndarray | MatrixExpr, /) -> MatrixExpr: ...
+    def __radd__(self, other: float, /) -> Expr: ...
+    @overload
+    def __sub__(self, other: float | Expr, /) -> Expr: ...
+    @overload
+    def __sub__(self, other: GenExpr, /) -> SumExpr: ...
+    @overload
+    def __sub__(self, other: np.ndarray | MatrixExpr, /) -> MatrixExpr: ...
+    def __rsub__(self, other: float, /) -> Expr: ...
+    @overload
+    def __mul__(self, other: float | Expr, /) -> Expr: ...
+    @overload
+    def __mul__(self, other: GenExpr, /) -> ProdExpr: ...
+    @overload
+    def __mul__(self, other: np.ndarray | MatrixExpr, /) -> MatrixExpr: ...
+    def __rmul__(self, other: float, /) -> Expr: ...
+    @overload
+    def __truediv__(self, other: float, /) -> Expr: ...
+    @overload
+    def __truediv__(self, other: Expr | GenExpr, /) -> ProdExpr: ...
+    @overload
+    def __truediv__(self, other: np.ndarray | MatrixExpr, /) -> MatrixExpr: ...
+    def __rtruediv__(self, other: float, /) -> ProdExpr: ...
+    # __pow__'s return type depends on the value of `other`:
+    # non-negative integers (including integer-valued floats) yield Expr,
+    # negative integers and non-integers yield PowExpr
+    # We annotate the int case as Expr to handle the common case (e.g. x**2)
+    # knowing it might be wrong in some cases, and likewise the float case as
+    # PowExpr can be wrong. We aim to make the stubs useful in common cases
+    # even though the type system cannot fully express the real types.
+    @overload
+    def __pow__(self, other: int, mod: Incomplete = ..., /) -> Expr: ...
+    @overload
+    def __pow__(self, other: float, mod: Incomplete = ..., /) -> PowExpr: ...
+    def __rpow__(self, other: float, /) -> UnaryExpr: ...
     def __getitem__(self, index: Incomplete, /) -> Incomplete: ...
-    def __gt__(self, other: object, /) -> bool: ...
-    def __iadd__(self, other: Incomplete, /) -> Incomplete: ...  # noqa: PYI034
     def __iter__(self) -> Incomplete: ...
-    def __le__(self, other: object, /) -> bool: ...
-    def __lt__(self, other: object, /) -> bool: ...
-    def __mul__(self, other: Incomplete, /) -> Incomplete: ...
-    def __ne__(self, other: object, /) -> bool: ...
-    def __pow__(self, other: Incomplete, modulo: Incomplete = ..., /) -> Incomplete: ...
-    def __rpow__(self, other: Incomplete, /) -> Incomplete: ...
-    def __truediv__(self, other: Incomplete, /) -> Incomplete: ...
 
 @disjoint_base
 class ExprCons:
@@ -380,12 +434,8 @@ class ExprCons:
     ) -> None: ...
     def normalize(self) -> Incomplete: ...
     def __bool__(self) -> bool: ...
-    def __eq__(self, other: object, /) -> bool: ...
-    def __ge__(self, other: object, /) -> bool: ...
-    def __gt__(self, other: object, /) -> bool: ...
-    def __le__(self, other: object, /) -> bool: ...
-    def __lt__(self, other: object, /) -> bool: ...
-    def __ne__(self, other: object, /) -> bool: ...
+    def __ge__(self, other: float, /) -> ExprCons: ...
+    def __le__(self, other: float, /) -> ExprCons: ...
 
 @disjoint_base
 class GenExpr(ExprLike):
@@ -394,17 +444,29 @@ class GenExpr(ExprLike):
     def __init__(self) -> None: ...
     def degree(self) -> Incomplete: ...
     def getOp(self) -> Incomplete: ...
-    def __add__(self, other: Incomplete, /) -> Incomplete: ...
-    def __eq__(self, other: object, /) -> bool: ...
-    def __ge__(self, other: object, /) -> bool: ...
-    def __gt__(self, other: object, /) -> bool: ...
-    def __le__(self, other: object, /) -> bool: ...
-    def __lt__(self, other: object, /) -> bool: ...
-    def __mul__(self, other: Incomplete, /) -> Incomplete: ...
-    def __ne__(self, other: object, /) -> bool: ...
-    def __pow__(self, other: Incomplete, modulo: Incomplete = ..., /) -> Incomplete: ...
-    def __rpow__(self, other: Incomplete, /) -> Incomplete: ...
-    def __truediv__(self, other: Incomplete, /) -> Incomplete: ...
+    def __neg__(self, /) -> ProdExpr: ...
+    @overload
+    def __add__(self, other: float | ExprLike, /) -> SumExpr: ...
+    @overload
+    def __add__(self, other: np.ndarray | MatrixExpr, /) -> MatrixExpr: ...
+    def __radd__(self, other: float, /) -> SumExpr: ...
+    @overload
+    def __sub__(self, other: float | ExprLike, /) -> SumExpr: ...
+    @overload
+    def __sub__(self, other: np.ndarray | MatrixExpr, /) -> MatrixExpr: ...
+    def __rsub__(self, other: float, /) -> SumExpr: ...
+    @overload
+    def __mul__(self, other: float | ExprLike, /) -> ProdExpr: ...
+    @overload
+    def __mul__(self, other: np.ndarray | MatrixExpr, /) -> MatrixExpr: ...
+    def __rmul__(self, other: float, /) -> ProdExpr: ...
+    @overload
+    def __truediv__(self, other: float | ExprLike, /) -> ProdExpr: ...
+    @overload
+    def __truediv__(self, other: np.ndarray | MatrixExpr, /) -> MatrixExpr: ...
+    def __rtruediv__(self, other: float, /) -> ProdExpr: ...
+    def __pow__(self, other: float | Constant, mod: Incomplete = ..., /) -> PowExpr: ...
+    def __rpow__(self, other: float, mod: Incomplete = ..., /) -> UnaryExpr: ...
 
 @disjoint_base
 class Heur:
@@ -534,6 +596,38 @@ class MatrixExpr(np.ndarray):
         *args: Incomplete,
         **kwargs: Incomplete,
     ) -> Incomplete: ...
+    def __neg__(self, /) -> MatrixExpr: ...
+    def __abs__(self, /) -> MatrixExpr: ...
+    def __add__(  # type: ignore[override]
+        self, other: float | ExprLike | np.ndarray | MatrixExpr, /
+    ) -> MatrixExpr: ...
+    def __radd__(self, other: float, /) -> MatrixExpr: ...  # type: ignore[override]
+    def __sub__(  # type: ignore[override]
+        self, other: float | ExprLike | np.ndarray | MatrixExpr, /
+    ) -> MatrixExpr: ...
+    def __rsub__(self, other: float, /) -> MatrixExpr: ...  # type: ignore[override]
+    def __mul__(  # type: ignore[override]
+        self, other: float | ExprLike | np.ndarray | MatrixExpr, /
+    ) -> MatrixExpr: ...
+    def __rmul__(self, other: float, /) -> MatrixExpr: ...  # type: ignore[override]
+    def __truediv__(  # type: ignore[override]
+        self, other: float | ExprLike | np.ndarray | MatrixExpr, /
+    ) -> MatrixExpr: ...
+    def __rtruediv__(self, other: float, /) -> MatrixExpr: ...  # type: ignore[override]
+    def __eq__(  # type: ignore[override]
+        self, other: float | ExprLike | np.ndarray | MatrixExpr, /
+    ) -> MatrixExprCons: ...
+    def __le__(  # type: ignore[override]
+        self, other: float | ExprLike | np.ndarray | MatrixExpr, /
+    ) -> MatrixExprCons: ...
+    def __ge__(  # type: ignore[override]
+        self, other: float | ExprLike | np.ndarray | MatrixExpr, /
+    ) -> MatrixExprCons: ...
+    def __matmul__(self, other: np.ndarray | MatrixExpr, /) -> MatrixExpr: ...  # type: ignore[override]
+    def __pow__(  # type: ignore[override]
+        self, other: float | np.ndarray, mod: Incomplete = ..., /
+    ) -> MatrixExpr: ...
+    def __rpow__(self, other: float, mod: Incomplete = ..., /) -> MatrixExpr: ...  # type: ignore[override]
 
 class MatrixExprCons(np.ndarray):
     def __array_ufunc__(
@@ -543,7 +637,8 @@ class MatrixExprCons(np.ndarray):
         *args: Incomplete,
         **kwargs: Incomplete,
     ) -> Incomplete: ...
-    def __eq__(self, other: object, /) -> bool: ...
+    def __ge__(self, other: float | np.ndarray, /) -> MatrixExprCons: ...  # type: ignore[override]
+    def __le__(self, other: float | np.ndarray, /) -> MatrixExprCons: ...  # type: ignore[override]
 
 class MatrixGenExpr(MatrixExpr): ...
 
@@ -562,6 +657,7 @@ class MatrixVariable(MatrixExpr):
     def isInLP(self) -> Incomplete: ...
     def varMayRound(self, direction: Incomplete = ...) -> Incomplete: ...
     def vtype(self) -> Incomplete: ...
+    def __pos__(self, /) -> MatrixExpr: ...  # type: ignore[override]
 
 @disjoint_base
 class Model:
@@ -2241,7 +2337,6 @@ class Term:
 
 class UnaryExpr(GenExpr):
     def __init__(self, *args: Incomplete, **kwargs: Incomplete) -> None: ...
-    def __abs__(self) -> GenExpr: ...
 
 @disjoint_base
 class VarExpr(GenExpr):

stubs/allowlist

@@ -1,3 +1,9 @@
 .*.__reduce_cython__
 .*.__setstate_cython__
 pyscipopt.scip.__test__
+# These exist at runtime but aren't useful in the stubs, they are added to the child classes
+pyscipopt.scip.ExprLike.__radd__
+pyscipopt.scip.ExprLike.__rmul__
+pyscipopt.scip.ExprLike.__rsub__
+pyscipopt.scip.ExprLike.__rtruediv__
+pyscipopt.scip.ExprLike.__sub__

stubs/baseline.sh

@@ -9,5 +9,8 @@ python "$REPO_ROOT"/scripts/generate_expr_type_tests.py
 for test_file in "$REPO_ROOT"/tests/@types/*.py; do
     echo "Updating mypy baseline for $test_file"
     output_file="${test_file%.*}.mypy.out"
-    python -m mypy "$test_file" --warn-unused-ignores | grep "error:" > "$output_file"
+    # write to a temp file first to avoid truncating the output file if mypy fails
+    tmp=$(mktemp)
+    python -m mypy "$test_file" --warn-unused-ignores | grep "error:" > "$tmp"
+    mv "$tmp" "$output_file"
 done