From d7aa31f3eb0971f8432694bd5e4d881dd1153866 Mon Sep 17 00:00:00 2001
From: "Alinson S. Xavier" <git@axavier.org>
Date: Tue, 13 Apr 2021 06:47:31 -0500
Subject: [PATCH] Fix mypy errors

---
 miplearn/components/static_lazy.py   | 8 +++++++-
 tests/components/test_static_lazy.py | 4 ++++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/miplearn/components/static_lazy.py b/miplearn/components/static_lazy.py
index bb65763..32b8dd5 100644
--- a/miplearn/components/static_lazy.py
+++ b/miplearn/components/static_lazy.py
@@ -60,6 +60,8 @@ class StaticLazyConstraintsComponent(Component):
         stats: LearningSolveStats,
         sample: Sample,
     ) -> None:
+        assert sample.after_mip is not None
+        assert sample.after_mip.extra is not None
         sample.after_mip.extra["lazy_enforced"] = self.enforced_cids
         stats["LazyStatic: Restored"] = self.n_restored
         stats["LazyStatic: Iterations"] = self.n_iterations
@@ -74,6 +76,10 @@ class StaticLazyConstraintsComponent(Component):
         sample: Sample,
     ) -> None:
         assert solver.internal_solver is not None
+        assert sample.after_load is not None
+        assert sample.after_load.instance is not None
+        assert sample.after_load.constraints is not None
+
         logger.info("Predicting violated (static) lazy constraints...")
         if sample.after_load.instance.lazy_constraint_count == 0:
             logger.info("Instance does not have static lazy constraints. Skipping.")
@@ -149,7 +155,7 @@ class StaticLazyConstraintsComponent(Component):
         _: Optional[Instance],
         sample: Sample,
     ) -> Tuple[Dict[Hashable, List[List[float]]], Dict[Hashable, List[List[float]]]]:
-        x, y, _ = self._sample_xy_with_cids(sample)
+        x, y, __ = self._sample_xy_with_cids(sample)
         return x, y
 
     def _check_and_add(self, solver: "LearningSolver") -> bool:
diff --git a/tests/components/test_static_lazy.py b/tests/components/test_static_lazy.py
index 047b578..5c3e6d5 100644
--- a/tests/components/test_static_lazy.py
+++ b/tests/components/test_static_lazy.py
@@ -114,6 +114,10 @@ def test_usage_with_solver(instance: Instance) -> None:
 
     stats: LearningSolveStats = {}
     sample = instance.samples[0]
+    assert sample.after_load is not None
+    assert sample.after_load.constraints is not None
+    assert sample.after_mip is not None
+    assert sample.after_mip.extra is not None
     del sample.after_mip.extra["lazy_enforced"]
 
     # LearningSolver calls before_solve_mip