[WIP] fix --check-untyped-defs for MyPy

pandas/core/computation/engines.py

@@ -47,14 +47,14 @@ def __init__(self, expr):
         self.aligned_axes = None
         self.result_type = None
 
-    def convert(self):
+    def convert(self) -> str:
         """Convert an expression for evaluation.
 
         Defaults to return the expression as a string.
         """
         return printing.pprint_thing(self.expr)
 
-    def evaluate(self):
+    def evaluate(self) -> object:
         """
         Run the engine on the expression.
 
@@ -76,7 +76,7 @@ def evaluate(self):
         )
 
     @property
-    def _is_aligned(self):
+    def _is_aligned(self) -> bool:
         return self.aligned_axes is not None and self.result_type is not None
 
     @abc.abstractmethod
@@ -102,12 +102,6 @@ class NumExprEngine(AbstractEngine):
 
     has_neg_frac = True
 
-    def __init__(self, expr):
-        super().__init__(expr)
-
-    def convert(self):
-        return str(super().convert())
-
     def _evaluate(self):
         import numexpr as ne
 
@@ -120,12 +114,8 @@ def _evaluate(self):
             _check_ne_builtin_clash(self.expr)
             return ne.evaluate(s, local_dict=scope)
         except KeyError as e:
-            # python 3 compat kludge
-            try:
-                msg = e.message
-            except AttributeError:
-                msg = str(e)
-            raise UndefinedVariableError(msg)
+            name = str(e)
+            raise UndefinedVariableError(name=name)
 
 
 class PythonEngine(AbstractEngine):
@@ -137,13 +127,10 @@ class PythonEngine(AbstractEngine):
 
     has_neg_frac = False
 
-    def __init__(self, expr):
-        super().__init__(expr)
-
     def evaluate(self):
         return self.expr()
 
-    def _evaluate(self):
+    def _evaluate(self) -> None:
         pass
 
 

pandas/core/computation/ops.py

@@ -5,7 +5,7 @@
 from distutils.version import LooseVersion
 from functools import partial
 import operator as op
-from typing import Callable, Iterable, Union
+from typing import Callable, Iterable, Optional, Union
 
 import numpy as np
 
@@ -56,7 +56,7 @@ class UndefinedVariableError(NameError):
     NameError subclass for local variables.
     """
 
-    def __init__(self, name, is_local):
+    def __init__(self, name: str, is_local: Optional[bool] = None):
         if is_local:
             msg = "local variable {0!r} is not defined"
         else:

pandas/core/computation/scope.py

@@ -170,7 +170,7 @@ def has_resolvers(self):
         """
         return bool(len(self.resolvers))
 
-    def resolve(self, key, is_local):
+    def resolve(self, key: str, is_local: bool) -> object:
         """
         Resolve a variable name in a possibly local context.
 

pandas/tests/frame/test_query_eval.py

@@ -481,11 +481,13 @@ def test_query_scope(self):
         assert_frame_equal(res, expected)
 
         # no local variable c
-        with pytest.raises(UndefinedVariableError):
+        with pytest.raises(
+            UndefinedVariableError, match="local variable 'c' is not defined"
+        ):
             df.query("@a > b > @c", engine=engine, parser=parser)
 
         # no column named 'c'
-        with pytest.raises(UndefinedVariableError):
+        with pytest.raises(UndefinedVariableError, match="name 'c' is not defined"):
             df.query("@a > b > c", engine=engine, parser=parser)
 
     def test_query_doesnt_pickup_local(self):
@@ -496,7 +498,7 @@ def test_query_doesnt_pickup_local(self):
         df = DataFrame(np.random.randint(m, size=(n, 3)), columns=list("abc"))
 
         # we don't pick up the local 'sin'
-        with pytest.raises(UndefinedVariableError):
+        with pytest.raises(UndefinedVariableError, match="name 'sin' is not defined"):
             df.query("sin > 5", engine=engine, parser=parser)
 
     def test_query_builtin(self):
@@ -590,7 +592,7 @@ def test_nested_raises_on_local_self_reference(self):
         df = DataFrame(np.random.randn(5, 3))
 
         # can't reference ourself b/c we're a local so @ is necessary
-        with pytest.raises(UndefinedVariableError):
+        with pytest.raises(UndefinedVariableError, match="name 'df' is not defined"):
             df.query("df > 0", engine=self.engine, parser=self.parser)
 
     def test_local_syntax(self):
@@ -653,9 +655,9 @@ def test_query_undefined_local(self):
         skip_if_no_pandas_parser(parser)
 
         df = DataFrame(np.random.rand(10, 2), columns=list("ab"))
-        msg = "local variable 'c' is not defined"
-
-        with pytest.raises(UndefinedVariableError, match=msg):
+        with pytest.raises(
+            UndefinedVariableError, match="local variable 'c' is not defined"
+        ):
             df.query("a == @c", engine=engine, parser=parser)
 
     def test_index_resolvers_come_after_columns_with_the_same_name(self):
@@ -787,7 +789,7 @@ def test_nested_scope(self):
         with pytest.raises(SyntaxError):
             df.query("(@df>0) & (@df2>0)", engine=engine, parser=parser)
 
-        with pytest.raises(UndefinedVariableError):
+        with pytest.raises(UndefinedVariableError, match="name 'df' is not defined"):
             df.query("(df>0) & (df2>0)", engine=engine, parser=parser)
 
         expected = df[(df > 0) & (df2 > 0)]

setup.cfg

@@ -191,9 +191,6 @@ check_untyped_defs=False
 [mypy-pandas.core.computation.align]
 check_untyped_defs=False
 
-[mypy-pandas.core.computation.engines]
-check_untyped_defs=False
-
 [mypy-pandas.core.computation.expr]
 check_untyped_defs=False