Add types to remaining files; activate mypy's disallow_untyped_defs

tests/problems/test_knapsack.py

@@ -8,7 +8,7 @@ from scipy.stats import uniform, randint
 from miplearn.problems.knapsack import MultiKnapsackGenerator
 
 
-def test_knapsack_generator():
+def test_knapsack_generator() -> None:
     gen = MultiKnapsackGenerator(
         n=randint(low=100, high=101),
         m=randint(low=30, high=31),
@@ -20,5 +20,5 @@ def test_knapsack_generator():
     instances = gen.generate(100)
     w_sum = sum(instance.weights for instance in instances) / len(instances)
     b_sum = sum(instance.capacities for instance in instances) / len(instances)
-    assert round(np.mean(w_sum), -1) == 500.0
-    assert round(np.mean(b_sum), -3) == 25000.0
+    assert round(float(np.mean(w_sum)), -1) == 500.0
+    assert round(float(np.mean(b_sum)), -3) == 25000.0

tests/problems/test_stab.py

@@ -10,16 +10,16 @@ from miplearn.problems.stab import MaxWeightStableSetInstance
 from miplearn.solvers.learning import LearningSolver
 
 
-def test_stab():
+def test_stab() -> None:
     graph = nx.cycle_graph(5)
-    weights = [1.0, 1.0, 1.0, 1.0, 1.0]
+    weights = np.array([1.0, 1.0, 1.0, 1.0, 1.0])
     instance = MaxWeightStableSetInstance(graph, weights)
     solver = LearningSolver()
     stats = solver.solve(instance)
     assert stats["Lower bound"] == 2.0
 
 
-def test_stab_generator_fixed_graph():
+def test_stab_generator_fixed_graph() -> None:
     np.random.seed(42)
     from miplearn.problems.stab import MaxWeightStableSetGenerator
 
@@ -36,7 +36,7 @@ def test_stab_generator_fixed_graph():
     assert list(weights_avg_actual) == weights_avg_expected
 
 
-def test_stab_generator_random_graph():
+def test_stab_generator_random_graph() -> None:
     np.random.seed(42)
     from miplearn.problems.stab import MaxWeightStableSetGenerator
 

tests/problems/test_tsp.py

@@ -11,7 +11,7 @@ from miplearn.problems.tsp import TravelingSalesmanGenerator, TravelingSalesmanI
 from miplearn.solvers.learning import LearningSolver
 
 
-def test_generator():
+def test_generator() -> None:
     instances = TravelingSalesmanGenerator(
         x=uniform(loc=0.0, scale=1000.0),
         y=uniform(loc=0.0, scale=1000.0),
@@ -26,7 +26,7 @@ def test_generator():
     assert np.std(d) > 0
 
 
-def test_instance():
+def test_instance() -> None:
     n_cities = 4
     distances = np.array(
         [
@@ -40,6 +40,7 @@ def test_instance():
     solver = LearningSolver()
     stats = solver.solve(instance)
     solution = instance.training_data[0].solution
+    assert solution is not None
     assert solution["x[(0, 1)]"] == 1.0
     assert solution["x[(0, 2)]"] == 0.0
     assert solution["x[(0, 3)]"] == 1.0
@@ -50,7 +51,7 @@ def test_instance():
     assert stats["Upper bound"] == 4.0
 
 
-def test_subtour():
+def test_subtour() -> None:
     n_cities = 6
     cities = np.array(
         [
@@ -66,8 +67,10 @@ def test_subtour():
     instance = TravelingSalesmanInstance(n_cities, distances)
     solver = LearningSolver()
     solver.solve(instance)
+    assert instance.training_data[0].lazy_enforced is not None
     assert len(instance.training_data[0].lazy_enforced) > 0
     solution = instance.training_data[0].solution
+    assert solution is not None
     assert solution["x[(0, 1)]"] == 1.0
     assert solution["x[(0, 4)]"] == 1.0
     assert solution["x[(1, 2)]"] == 1.0