fix typing errors

Author: coroa

linopy/common.py

@@ -19,6 +19,7 @@
 import pandas as pd
 import polars as pl
 from numpy import arange, signedinteger
+from polars.datatypes import DataTypeClass
 from xarray import DataArray, Dataset, apply_ufunc, broadcast
 from xarray import align as xr_align
 from xarray.core import dtypes, indexing
@@ -327,7 +328,7 @@ def check_has_nulls(df: pd.DataFrame, name: str) -> None:
         raise ValueError(f"Fields {name} contains nan's in field(s) {fields}")
 
 
-def infer_schema_polars(ds: Dataset) -> dict[Hashable, pl.DataType]:
+def infer_schema_polars(ds: Dataset) -> dict[str, DataTypeClass]:
     """
     Infer the polars data schema from a xarray dataset.
 
@@ -339,21 +340,22 @@ def infer_schema_polars(ds: Dataset) -> dict[Hashable, pl.DataType]:
     -------
         dict: A dictionary mapping column names to their corresponding Polars data types.
     """
-    schema = {}
+    schema: dict[str, DataTypeClass] = {}
     np_major_version = int(np.__version__.split(".")[0])
     use_int32 = os.name == "nt" and np_major_version < 2
     for name, array in ds.items():
+        name = str(name)
         if np.issubdtype(array.dtype, np.integer):
             schema[name] = pl.Int32 if use_int32 else pl.Int64
         elif np.issubdtype(array.dtype, np.floating):
-            schema[name] = pl.Float64  # type: ignore
+            schema[name] = pl.Float64
         elif np.issubdtype(array.dtype, np.bool_):
-            schema[name] = pl.Boolean  # type: ignore
+            schema[name] = pl.Boolean
         elif np.issubdtype(array.dtype, np.object_):
-            schema[name] = pl.Object  # type: ignore
+            schema[name] = pl.Object
         else:
-            schema[name] = pl.Utf8  # type: ignore
-    return schema  # type: ignore
+            schema[name] = pl.Utf8
+    return schema
 
 
 def to_polars(ds: Dataset, **kwargs: Any) -> pl.DataFrame:
@@ -429,7 +431,7 @@ def filter_nulls_polars(df: pl.DataFrame) -> pl.DataFrame:
     if "labels" in df.columns:
         cond.append(pl.col("labels").ne(-1))
 
-    cond = reduce(operator.and_, cond)  # type: ignore
+    cond = reduce(operator.and_, cond)  # type: ignore[arg-type]
     return df.filter(cond)
 
 
@@ -655,7 +657,7 @@ def iterate_slices(
         start = i * chunk_size
         end = min(start + chunk_size, size_of_leading_dim)
         slice_dict = {leading_dim: slice(start, end)}
-        yield ds.isel(slice_dict)
+        yield ds.isel(slice_dict)  # type: ignore[attr-defined]
 
 
 def _remap(array: np.ndarray, mapping: np.ndarray) -> np.ndarray:
@@ -1367,7 +1369,7 @@ def __getitem__(
             # expand the indexer so we can handle Ellipsis
             labels = indexing.expanded_indexer(key, self.object.ndim)
             key = dict(zip(self.object.dims, labels))
-        return self.object.sel(key)
+        return self.object.sel(key)  # type: ignore[attr-defined]
 
 
 class EmptyDeprecationWrapper:
@@ -1439,9 +1441,9 @@ def coords_from_dataset(ds: Dataset, coord_dims: list[str]) -> list[pd.Index]:
         if f"_coord_{d}_codes" in ds:
             codes_2d = ds[f"_coord_{d}_codes"].values.T
             level_names = [
-                k[len(f"_coord_{d}_level_") :]
+                str(k)[len(f"_coord_{d}_level_") :]
                 for k in ds
-                if k.startswith(f"_coord_{d}_level_")
+                if str(k).startswith(f"_coord_{d}_level_")
             ]
             arrays = [
                 ds[f"_coord_{d}_level_{ln}"].values[codes_2d[i]]

linopy/constants.py

@@ -214,9 +214,11 @@ def process(cls, status: str, termination_condition: str) -> "Status":
 
     @classmethod
     def from_termination_condition(
-        cls, termination_condition: Union["TerminationCondition", str]
+        cls, termination_condition: Union["TerminationCondition", str, None]
     ) -> "Status":
-        termination_condition = TerminationCondition.process(termination_condition)
+        termination_condition = TerminationCondition.process(
+            termination_condition if termination_condition is not None else "unknown"
+        )
         solver_status = SolverStatus.from_termination_condition(termination_condition)
         return cls(solver_status, termination_condition)
 

linopy/constraints.py

@@ -169,6 +169,11 @@ def rhs(self) -> DataArray:
     def dual(self) -> DataArray:
         """Get the dual values DataArray."""
 
+    @dual.setter
+    @abstractmethod
+    def dual(self, value: DataArray) -> None:
+        """Set the dual values DataArray."""
+
     @abstractmethod
     def has_variable(self, variable: variables.Variable) -> bool:
         """Check if the constraint references any of the given variable labels."""
@@ -185,6 +190,18 @@ def sanitize_missings(self) -> ConstraintBase:
     def sanitize_infinities(self) -> ConstraintBase:
         """Mask out rows with invalid infinite RHS values."""
 
+    @abstractmethod
+    def to_polars(self) -> pl.DataFrame:
+        """Convert constraint to a polars DataFrame."""
+
+    @abstractmethod
+    def freeze(self) -> Constraint:
+        """Return an immutable Constraint (CSR-backed)."""
+
+    @abstractmethod
+    def mutable(self) -> MutableConstraint:
+        """Return a mutable MutableConstraint."""
+
     @abstractmethod
     def to_matrix_with_rhs(
         self, label_index: VariableLabelIndex
@@ -298,7 +315,8 @@ def mask(self) -> DataArray | None:
         (True) and disabled (False).
         """
         if self.is_assigned:
-            return (self.labels != FILL_VALUE["labels"]).astype(bool)
+            result: DataArray = self.labels != FILL_VALUE["labels"]  # type: ignore[assignment]
+            return result.astype(bool)
         return None
 
     @property
@@ -391,7 +409,7 @@ def flat(self) -> pd.DataFrame:
         """
         ds = self.data
 
-        def mask_func(data: pd.DataFrame) -> pd.Series:
+        def mask_func(data: dict) -> pd.Series:
             mask = (data["vars"] != -1) & (data["coeffs"] != 0)
             if "labels" in data:
                 mask &= data["labels"] != -1
@@ -593,7 +611,7 @@ def nterm(self) -> int:
 
     @property
     def coord_names(self) -> list[str]:
-        return [c.name for c in self._coords]
+        return [str(c.name) for c in self._coords]
 
     @property
     def labels(self) -> DataArray:
@@ -828,8 +846,8 @@ def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> Constraint:
         )
         rhs = ds["rhs"].values
         sign = attrs["sign"]
-        cindex = int(attrs["cindex"])
-        cindex = cindex if cindex >= 0 else None
+        _cindex_raw = int(attrs["cindex"])
+        cindex: int | None = _cindex_raw if _cindex_raw >= 0 else None
         coord_dims = attrs["coord_dims"]
         if isinstance(coord_dims, str):
             coord_dims = [coord_dims]
@@ -892,7 +910,7 @@ def mutable(self) -> MutableConstraint:
         """Convert to a MutableConstraint."""
         return MutableConstraint(self.data, self._model, self._name)
 
-    def to_polars(self) -> Any:
+    def to_polars(self) -> pl.DataFrame:
         """Convert to polars DataFrame — delegates to mutable()."""
         return self.mutable().to_polars()
 
@@ -1598,6 +1616,8 @@ def set_blocks(self, block_map: np.ndarray) -> None:
         N = block_map.max()
 
         for name, constraint in self.items():
+            if not isinstance(constraint, MutableConstraint):
+                self.data[name] = constraint = constraint.mutable()
             res = xr.full_like(constraint.labels, N + 1, dtype=block_map.dtype)
             entries = replace_by_map(constraint.vars, block_map)
 
@@ -1679,9 +1699,12 @@ def reset_dual(self) -> None:
                         cindex=c._cindex,
                         dual=None,
                     )
-            else:
+            elif isinstance(c, MutableConstraint):
                 if "dual" in c.data:
                     c._data = c.data.drop_vars("dual")
+            else:
+                msg = f"reset_dual encountered an unknown constraint type: {type(c)}"
+                raise NotImplementedError(msg)
 
 
 class AnonymousScalarConstraint: