Skip to content

Commit

Permalink
add typehints and ignore types to test to satisfy mypy
Browse files Browse the repository at this point in the history
  • Loading branch information
FBumann committed Feb 18, 2025
1 parent 615fb4d commit 084d112
Showing 1 changed file with 27 additions and 25 deletions.
52 changes: 27 additions & 25 deletions test/test_compatible_arithmetrics.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
from typing import Any

import numpy as np
import pandas as pd
import pytest
import xarray as xr

from linopy import Model
from linopy import LESS_EQUAL, Model, Variable
from linopy.testing import assert_linequal


Expand All @@ -26,40 +28,40 @@ def activate(self, active: int) -> None:
def active_data(self) -> xr.DataArray:
return self.data1 if self.active == 1 else self.data2

def __add__(self, other):
def __add__(self, other: Any) -> xr.DataArray:
return self.active_data + other

def __sub__(self, other):
def __sub__(self, other: Any) -> xr.DataArray:
return self.active_data - other

def __mul__(self, other):
def __mul__(self, other: Any) -> xr.DataArray:
return self.active_data * other

def __truediv__(self, other):
def __truediv__(self, other: Any) -> xr.DataArray:
return self.active_data / other

def __radd__(self, other):
def __radd__(self, other: Any) -> Any:
return other + self.active_data

def __rsub__(self, other):
def __rsub__(self, other: Any) -> Any:
return other - self.active_data

def __rmul__(self, other):
def __rmul__(self, other: Any) -> Any:
return other * self.active_data

def __rtruediv__(self, other):
def __rtruediv__(self, other: Any) -> Any:
return other / self.active_data

def __neg__(self):
def __neg__(self) -> xr.DataArray:
return -self.active_data

def __pos__(self):
def __pos__(self) -> xr.DataArray:
return +self.active_data

def __abs__(self):
def __abs__(self) -> xr.DataArray:
return abs(self.active_data)

def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): # type: ignore
# Ensure we always use the active_data when interacting with numpy/xarray operations
new_inputs = [
inp.active_data if isinstance(inp, SomeOtherDatatype) else inp
Expand All @@ -79,23 +81,23 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
],
ids=["single_dim", "multi_dim"],
)
def m(request) -> Model:
def m(request) -> Model: # type: ignore
m = Model()
m.add_variables(coords=request.param, name="x")
x = m.add_variables(coords=request.param, name="x")
m.add_variables(0, 10, name="z")
m.add_constraints(m.variables["x"] >= 0, name="c")
m.add_constraints(x, LESS_EQUAL, 0, name="c")
return m


def test_arithmetric_operations_variable(m: Model) -> None:
x = m.variables["x"]
x: Variable = m.variables["x"]
rng = np.random.default_rng()
data = xr.DataArray(rng.random(x.shape), coords=x.coords)
other_datatype = SomeOtherDatatype(data.copy())
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(x - data, x - other_datatype) # type: ignore
assert_linequal(x * data, x * other_datatype) # type: ignore
assert_linequal(x / data, x / other_datatype) # type: ignore


def test_arithmetric_operations_con(m: Model) -> None:
Expand All @@ -108,7 +110,7 @@ def test_arithmetric_operations_con(m: Model) -> None:
assert_linequal(c.lhs - data, c.lhs - other_datatype)
assert_linequal(c.lhs * data, c.lhs * other_datatype)
assert_linequal(c.lhs / data, c.lhs / other_datatype)
assert_linequal(c.rhs + data, c.rhs + other_datatype)
assert_linequal(c.rhs - data, c.rhs - other_datatype)
assert_linequal(c.rhs * data, c.rhs * other_datatype)
assert_linequal(c.rhs / data, c.rhs / other_datatype)
assert_linequal(c.rhs + data, c.rhs + other_datatype) # type: ignore
assert_linequal(c.rhs - data, c.rhs - other_datatype) # type: ignore
assert_linequal(c.rhs * data, c.rhs * other_datatype) # type: ignore
assert_linequal(c.rhs / data, c.rhs / other_datatype) # type: ignore

0 comments on commit 084d112

Please sign in to comment.