Fix mypy errors

miplearn/collectors/basic.py

@@ -21,7 +21,7 @@ class BasicCollector:
         n_jobs: int = 1,
         progress: bool = False,
     ) -> None:
-        def _collect(data_filename):
+        def _collect(data_filename: str) -> None:
             h5_filename = _to_h5_filename(data_filename)
             mps_filename = h5_filename.replace(".h5", ".mps")
 

miplearn/extractors/AlvLouWeh2017.py

@@ -22,7 +22,7 @@ class AlvLouWeh2017Extractor(FeaturesExtractor):
         self.with_m3 = with_m3
 
     def get_instance_features(self, h5: H5File) -> np.ndarray:
-        raise NotImplemented()
+        raise NotImplementedError()
 
     def get_var_features(self, h5: H5File) -> np.ndarray:
         """
@@ -197,7 +197,7 @@ class AlvLouWeh2017Extractor(FeaturesExtractor):
         return features
 
     def get_constr_features(self, h5: H5File) -> np.ndarray:
-        raise NotImplemented()
+        raise NotImplementedError()
 
 
 def _fix_infinity(m: Optional[np.ndarray]) -> None:

miplearn/extractors/fields.py

@@ -31,9 +31,9 @@ class H5FieldsExtractor(FeaturesExtractor):
                 data = h5.get_scalar(field)
             assert data is not None
             x.append(data)
-        x = np.hstack(x)
-        assert len(x.shape) == 1
-        return x
+        x_np = np.hstack(x)
+        assert len(x_np.shape) == 1
+        return x_np
 
     def get_var_features(self, h5: H5File) -> np.ndarray:
         var_types = h5.get_array("static_var_types")
@@ -51,13 +51,14 @@ class H5FieldsExtractor(FeaturesExtractor):
             raise Exception("No constr fields provided")
         return self._extract(h5, self.constr_fields, n_constr)
 
-    def _extract(self, h5, fields, n_expected):
+    def _extract(self, h5: H5File, fields: List[str], n_expected: int) -> np.ndarray:
         x = []
         for field in fields:
             try:
                 data = h5.get_array(field)
             except ValueError:
                 v = h5.get_scalar(field)
+                assert v is not None
                 data = np.repeat(v, n_expected)
             assert data is not None
             assert len(data.shape) == 1

miplearn/h5.py

@@ -111,7 +111,7 @@ class H5File:
         ), f"bytes expected; found: {value.__class__}"  # type: ignore
         self.put_array(key, np.frombuffer(value, dtype="uint8"))
 
-    def close(self):
+    def close(self) -> None:
         self.file.close()
 
     def __enter__(self) -> "H5File":

miplearn/problems/setcover.py

@@ -95,7 +95,7 @@ def build_setcover_model_gurobipy(data: Union[str, SetCoverData]) -> GurobiModel
 
 def build_setcover_model_pyomo(
     data: Union[str, SetCoverData],
-    solver="gurobi_persistent",
+    solver: str = "gurobi_persistent",
 ) -> PyomoModel:
     data = _read_setcover_data(data)
     (n_elements, n_sets) = data.incidence_matrix.shape

miplearn/problems/stab.py

@@ -96,7 +96,7 @@ def build_stab_model_gurobipy(data: MaxWeightStableSetData) -> GurobiModel:
 
 def build_stab_model_pyomo(
     data: MaxWeightStableSetData,
-    solver="gurobi_persistent",
+    solver: str = "gurobi_persistent",
 ) -> PyomoModel:
     data = _read_stab_data(data)
     model = pe.ConcreteModel()

miplearn/solvers/gurobi.py

@@ -9,9 +9,10 @@ import numpy as np
 from scipy.sparse import lil_matrix
 
 from miplearn.h5 import H5File
+from miplearn.solvers.abstract import AbstractModel
 
 
-class GurobiModel:
+class GurobiModel(AbstractModel):
     _supports_basis_status = True
     _supports_sensitivity_analysis = True
     _supports_node_count = True

miplearn/solvers/learning.py

@@ -3,7 +3,7 @@
 #  Released under the modified BSD license. See COPYING.md for more details.
 from os.path import exists
 from tempfile import NamedTemporaryFile
-from typing import List, Any, Union
+from typing import List, Any, Union, Dict, Callable, Optional
 
 from miplearn.h5 import H5File
 from miplearn.io import _to_h5_filename
@@ -11,23 +11,28 @@ from miplearn.solvers.abstract import AbstractModel
 
 
 class LearningSolver:
-    def __init__(self, components: List[Any], skip_lp=False):
+    def __init__(self, components: List[Any], skip_lp: bool = False) -> None:
         self.components = components
         self.skip_lp = skip_lp
 
-    def fit(self, data_filenames):
+    def fit(self, data_filenames: List[str]) -> None:
         h5_filenames = [_to_h5_filename(f) for f in data_filenames]
         for comp in self.components:
             comp.fit(h5_filenames)
 
-    def optimize(self, model: Union[str, AbstractModel], build_model=None):
+    def optimize(
+        self,
+        model: Union[str, AbstractModel],
+        build_model: Optional[Callable] = None,
+    ) -> Dict[str, Any]:
         if isinstance(model, str):
             h5_filename = _to_h5_filename(model)
             assert build_model is not None
             model = build_model(model)
+            assert isinstance(model, AbstractModel)
         else:
             h5_filename = NamedTemporaryFile().name
-        stats = {}
+        stats: Dict[str, Any] = {}
         mode = "r+" if exists(h5_filename) else "w"
         with H5File(h5_filename, mode) as h5:
             model.extract_after_load(h5)

miplearn/solvers/pyomo.py

@@ -2,7 +2,7 @@
 #  Copyright (C) 2020-2022, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 from numbers import Number
-from typing import Optional, Dict, List, Any
+from typing import Optional, Dict, List, Any, Tuple, Union
 
 import numpy as np
 import pyomo
@@ -24,7 +24,7 @@ class PyomoModel(AbstractModel):
         self.is_persistent = hasattr(self.solver, "set_instance")
         if self.is_persistent:
             self.solver.set_instance(model)
-        self.results = None
+        self.results: Optional[Dict] = None
         self._is_warm_start_available = False
         if not hasattr(self.inner, "dual"):
             self.inner.dual = Suffix(direction=Suffix.IMPORT)
@@ -56,7 +56,7 @@ class PyomoModel(AbstractModel):
                 raise Exception(f"Unknown sense: {sense}")
             self.solver.add_constraint(eq)
 
-    def _var_names_to_vars(self, var_names):
+    def _var_names_to_vars(self, var_names: np.ndarray) -> List[Any]:
         varname_to_var = {}
         for var in self.inner.component_objects(Var):
             for idx in var:
@@ -70,12 +70,14 @@ class PyomoModel(AbstractModel):
         h5.put_scalar("static_sense", self._get_sense())
 
     def extract_after_lp(self, h5: H5File) -> None:
+        assert self.results is not None
         self._extract_after_lp_vars(h5)
         self._extract_after_lp_constrs(h5)
         h5.put_scalar("lp_obj_value", self.results["Problem"][0]["Lower bound"])
         h5.put_scalar("lp_wallclock_time", self._get_runtime())
 
-    def _get_runtime(self):
+    def _get_runtime(self) -> float:
+        assert self.results is not None
         solver_dict = self.results["Solver"][0]
         for key in ["Wallclock time", "User time"]:
             if isinstance(solver_dict[key], Number):
@@ -83,6 +85,7 @@ class PyomoModel(AbstractModel):
         raise Exception("Time unavailable")
 
     def extract_after_mip(self, h5: H5File) -> None:
+        assert self.results is not None
         h5.put_scalar("mip_wallclock_time", self._get_runtime())
         if self.results["Solver"][0]["Termination condition"] == "infeasible":
             return
@@ -150,7 +153,7 @@ class PyomoModel(AbstractModel):
                 var.value = val
         self._is_warm_start_available = True
 
-    def _extract_after_load_vars(self, h5):
+    def _extract_after_load_vars(self, h5: H5File) -> None:
         names: List[str] = []
         types: List[str] = []
         upper_bounds: List[float] = []
@@ -211,7 +214,7 @@ class PyomoModel(AbstractModel):
         h5.put_array("static_var_obj_coeffs", np.array(obj_coeffs))
         h5.put_scalar("static_obj_offset", obj_offset)
 
-    def _extract_after_load_constrs(self, h5):
+    def _extract_after_load_constrs(self, h5: H5File) -> None:
         names: List[str] = []
         rhs: List[float] = []
         senses: List[str] = []
@@ -219,7 +222,7 @@ class PyomoModel(AbstractModel):
         lhs_col: List[int] = []
         lhs_data: List[float] = []
 
-        varname_to_idx = {}
+        varname_to_idx: Dict[str, int] = {}
         for var in self.inner.component_objects(Var):
             for idx in var:
                 varname = var.name
@@ -285,7 +288,7 @@ class PyomoModel(AbstractModel):
         h5.put_array("static_constr_rhs", np.array(rhs))
         h5.put_array("static_constr_sense", np.array(senses, dtype="S"))
 
-    def _extract_after_lp_vars(self, h5):
+    def _extract_after_lp_vars(self, h5: H5File) -> None:
         rc = []
         values = []
         for var in self.inner.component_objects(Var):
@@ -296,7 +299,7 @@ class PyomoModel(AbstractModel):
         h5.put_array("lp_var_reduced_costs", np.array(rc))
         h5.put_array("lp_var_values", np.array(values))
 
-    def _extract_after_lp_constrs(self, h5):
+    def _extract_after_lp_constrs(self, h5: H5File) -> None:
         dual = []
         slacks = []
         for constr in self.inner.component_objects(pyomo.core.Constraint):
@@ -307,7 +310,7 @@ class PyomoModel(AbstractModel):
         h5.put_array("lp_constr_dual_values", np.array(dual))
         h5.put_array("lp_constr_slacks", np.array(slacks))
 
-    def _extract_after_mip_vars(self, h5):
+    def _extract_after_mip_vars(self, h5: H5File) -> None:
         values = []
         for var in self.inner.component_objects(Var):
             for idx in var:
@@ -315,7 +318,7 @@ class PyomoModel(AbstractModel):
                 values.append(v.value)
         h5.put_array("mip_var_values", np.array(values))
 
-    def _extract_after_mip_constrs(self, h5):
+    def _extract_after_mip_constrs(self, h5: H5File) -> None:
         slacks = []
         for constr in self.inner.component_objects(pyomo.core.Constraint):
             for idx in constr:
@@ -323,7 +326,7 @@ class PyomoModel(AbstractModel):
                 slacks.append(abs(self.inner.slack[c]))
         h5.put_array("mip_constr_slacks", np.array(slacks))
 
-    def _parse_pyomo_expr(self, expr: Any):
+    def _parse_pyomo_expr(self, expr: Any) -> Tuple[Dict[str, float], float]:
         lhs = {}
         offset = 0.0
         if isinstance(expr, SumExpression):
@@ -332,7 +335,7 @@ class PyomoModel(AbstractModel):
                     lhs[term._args_[1].name] = float(term._args_[0])
                 elif isinstance(term, _GeneralVarData):
                     lhs[term.name] = 1.0
-                elif isinstance(term, Number):
+                elif isinstance(term, float):
                     offset += term
                 else:
                     raise Exception(f"Unknown term type: {term.__class__.__name__}")
@@ -342,7 +345,7 @@ class PyomoModel(AbstractModel):
             raise Exception(f"Unknown expression type: {expr.__class__.__name__}")
         return lhs, offset
 
-    def _gap(self, zp, zd, tol=1e-6):
+    def _gap(self, zp: float, zd: float, tol: float = 1e-6) -> float:
         # Reference: https://www.gurobi.com/documentation/9.5/refman/mipgap2.html
         if abs(zp) < tol:
             if abs(zd) < tol:
@@ -352,7 +355,7 @@ class PyomoModel(AbstractModel):
         else:
             return abs(zp - zd) / abs(zp)
 
-    def _get_sense(self):
+    def _get_sense(self) -> str:
         for obj in self.inner.component_objects(Objective):
             sense = obj.sense
             if sense == pyomo.core.kernel.objective.minimize:
@@ -361,6 +364,7 @@ class PyomoModel(AbstractModel):
                 return "max"
             else:
                 raise Exception(f"Unknown sense: ${sense}")
+        raise Exception(f"No objective")
 
     def write(self, filename: str) -> None:
         self.inner.write(filename, io_options={"symbolic_solver_labels": True})

tests/problems/test_setcover.py

@@ -14,6 +14,7 @@ from miplearn.problems.setcover import (
     SetCoverGenerator,
     build_setcover_model_pyomo,
 )
+from miplearn.solvers.abstract import AbstractModel
 
 
 def test_set_cover_generator() -> None:
@@ -84,6 +85,7 @@ def test_set_cover() -> None:
         build_setcover_model_pyomo(data),
         build_setcover_model_gurobipy(data),
     ]:
+        assert isinstance(model, AbstractModel)
         with NamedTemporaryFile() as tempfile:
             with H5File(tempfile.name) as h5:
                 model.optimize()

tests/problems/test_stab.py

@@ -12,6 +12,7 @@ from miplearn.problems.stab import (
     build_stab_model_pyomo,
     build_stab_model_gurobipy,
 )
+from miplearn.solvers.abstract import AbstractModel
 
 
 def test_stab() -> None:
@@ -23,6 +24,7 @@ def test_stab() -> None:
         build_stab_model_pyomo(data),
         build_stab_model_gurobipy(data),
     ]:
+        assert isinstance(model, AbstractModel)
         with NamedTemporaryFile() as tempfile:
             with H5File(tempfile.name) as h5:
                 model.optimize()

tests/test_solvers.py

@@ -3,6 +3,7 @@
 #  Released under the modified BSD license. See COPYING.md for more details.
 
 from tempfile import NamedTemporaryFile
+from typing import Callable, Any
 
 import numpy as np
 import pytest
@@ -40,28 +41,28 @@ def test_pyomo_persistent(data: SetCoverData) -> None:
     _test_solver(lambda d: build_setcover_model_pyomo(d, "gurobi_persistent"), data)
 
 
-def _test_solver(build_model, data):
+def _test_solver(build_model: Callable, data: Any) -> None:
     _test_extract(build_model(data))
     _test_add_constr(build_model(data))
     _test_fix_vars(build_model(data))
     _test_infeasible(build_model(data))
 
 
-def _test_extract(model):
+def _test_extract(model: AbstractModel) -> None:
     with NamedTemporaryFile() as tempfile:
         with H5File(tempfile.name) as h5:
 
-            def test_scalar(key, expected_value):
+            def test_scalar(key: str, expected_value: Any) -> None:
                 actual_value = h5.get_scalar(key)
                 assert actual_value is not None
                 assert actual_value == expected_value
 
-            def test_array(key, expected_value):
+            def test_array(key: str, expected_value: Any) -> None:
                 actual_value = h5.get_array(key)
                 assert actual_value is not None
                 assert actual_value.tolist() == expected_value
 
-            def test_sparse(key, expected_value):
+            def test_sparse(key: str, expected_value: Any) -> None:
                 actual_value = h5.get_sparse(key)
                 assert actual_value is not None
                 assert actual_value.todense().tolist() == expected_value
@@ -143,7 +144,7 @@ def _test_extract(model):
                 assert pool_var_values.shape == (n_sols, 5)
 
 
-def _test_add_constr(model: AbstractModel):
+def _test_add_constr(model: AbstractModel) -> None:
     with NamedTemporaryFile() as tempfile:
         with H5File(tempfile.name) as h5:
             model.add_constrs(
@@ -154,10 +155,12 @@ def _test_add_constr(model: AbstractModel):
             )
             model.optimize()
             model.extract_after_mip(h5)
-            assert h5.get_array("mip_var_values").tolist() == [1, 0, 0, 0, 1]
+            mip_var_values = h5.get_array("mip_var_values")
+            assert mip_var_values is not None
+            assert mip_var_values.tolist() == [1, 0, 0, 0, 1]
 
 
-def _test_fix_vars(model: AbstractModel):
+def _test_fix_vars(model: AbstractModel) -> None:
     with NamedTemporaryFile() as tempfile:
         with H5File(tempfile.name) as h5:
             model.fix_variables(
@@ -166,10 +169,12 @@ def _test_fix_vars(model: AbstractModel):
             )
             model.optimize()
             model.extract_after_mip(h5)
-            assert h5.get_array("mip_var_values").tolist() == [1, 0, 0, 0, 1]
+            mip_var_values = h5.get_array("mip_var_values")
+            assert mip_var_values is not None
+            assert mip_var_values.tolist() == [1, 0, 0, 0, 1]
 
 
-def _test_infeasible(model: AbstractModel):
+def _test_infeasible(model: AbstractModel) -> None:
     with NamedTemporaryFile() as tempfile:
         with H5File(tempfile.name) as h5:
             model.fix_variables(