From 2d4ded19787fa651fbcee402ef819fdd6f9f7cd2 Mon Sep 17 00:00:00 2001
From: "Alinson S. Xavier" <git@axavier.org>
Date: Sun, 11 Apr 2021 17:07:45 -0500
Subject: [PATCH] Fix some mypy issues

---
 miplearn/solvers/gurobi.py     | 2 +-
 miplearn/solvers/pyomo/base.py | 4 +++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/miplearn/solvers/gurobi.py b/miplearn/solvers/gurobi.py
index 7060c2a..7f9152f 100644
--- a/miplearn/solvers/gurobi.py
+++ b/miplearn/solvers/gurobi.py
@@ -464,7 +464,7 @@ class GurobiSolver(InternalSolver):
         return var
 
     @staticmethod
-    def _parse_gurobi_var_lp(gp_var, var):
+    def _parse_gurobi_var_lp(gp_var: Any, var: Variable) -> None:
         var.reduced_cost = gp_var.rc
         var.sa_obj_up = gp_var.saobjUp
         var.sa_obj_down = gp_var.saobjLow
diff --git a/miplearn/solvers/pyomo/base.py b/miplearn/solvers/pyomo/base.py
index 00a74ca..b7e231a 100644
--- a/miplearn/solvers/pyomo/base.py
+++ b/miplearn/solvers/pyomo/base.py
@@ -399,6 +399,8 @@ class BasePyomoSolver(InternalSolver):
         return None
 
     def _parse_pyomo_variable(self, var: pyomo.core.Var) -> Variable:
+        assert self.model is not None
+
         # Variable type
         vtype: Optional[str] = None
         if var.domain == pyomo.core.Binary:
@@ -472,7 +474,7 @@ class BasePyomoSolver(InternalSolver):
         # Build constraint
         return constr
 
-    def _parse_pyomo_expr(self, expr):
+    def _parse_pyomo_expr(self, expr: Any) -> Dict[str, float]:
         lhs = {}
         if isinstance(expr, SumExpression):
             for term in expr._args_: