TYP: annotate queryables (#29754)

Author: jbrockmendel

pandas/core/computation/ops.py

@@ -55,7 +55,7 @@ class UndefinedVariableError(NameError):
     NameError subclass for local variables.
     """
 
-    def __init__(self, name, is_local):
+    def __init__(self, name, is_local: bool):
         if is_local:
             msg = "local variable {0!r} is not defined"
         else:
@@ -69,6 +69,8 @@ def __new__(cls, name, env, side=None, encoding=None):
         supr_new = super(Term, klass).__new__
         return supr_new(klass)
 
+    is_local: bool
+
     def __init__(self, name, env, side=None, encoding=None):
         # name is a str for Term, but may be something else for subclasses
         self._name = name

pandas/core/computation/pytables.py

@@ -2,7 +2,7 @@
 
 import ast
 from functools import partial
-from typing import Any, Optional, Tuple
+from typing import Any, Dict, Optional, Tuple
 
 import numpy as np
 
@@ -24,17 +24,27 @@
 class Scope(_scope.Scope):
     __slots__ = ("queryables",)
 
-    def __init__(self, level: int, global_dict=None, local_dict=None, queryables=None):
+    queryables: Dict[str, Any]
+
+    def __init__(
+        self,
+        level: int,
+        global_dict=None,
+        local_dict=None,
+        queryables: Optional[Dict[str, Any]] = None,
+    ):
         super().__init__(level + 1, global_dict=global_dict, local_dict=local_dict)
         self.queryables = queryables or dict()
 
 
 class Term(ops.Term):
+    env: Scope
+
     def __new__(cls, name, env, side=None, encoding=None):
         klass = Constant if not isinstance(name, str) else cls
         return object.__new__(klass)
 
-    def __init__(self, name, env, side=None, encoding=None):
+    def __init__(self, name, env: Scope, side=None, encoding=None):
         super().__init__(name, env, side=side, encoding=encoding)
 
     def _resolve_name(self):
@@ -69,7 +79,10 @@ class BinOp(ops.BinOp):
 
     _max_selectors = 31
 
-    def __init__(self, op, lhs, rhs, queryables, encoding):
+    op: str
+    queryables: Dict[str, Any]
+
+    def __init__(self, op: str, lhs, rhs, queryables: Dict[str, Any], encoding):
         super().__init__(op, lhs, rhs)
         self.queryables = queryables
         self.encoding = encoding
@@ -373,9 +386,6 @@ def prune(self, klass):
         return None
 
 
-_op_classes = {"unary": UnaryOp}
-
-
 class ExprVisitor(BaseExprVisitor):
     const_type = Constant
     term_type = Term
@@ -510,7 +520,16 @@ class Expr(expr.Expr):
     "major_axis>=20130101"
     """
 
-    def __init__(self, where, queryables=None, encoding=None, scope_level: int = 0):
+    _visitor: Optional[ExprVisitor]
+    env: Scope
+
+    def __init__(
+        self,
+        where,
+        queryables: Optional[Dict[str, Any]] = None,
+        encoding=None,
+        scope_level: int = 0,
+    ):
 
         where = _validate_where(where)
 

pandas/core/computation/scope.py

@@ -9,6 +9,7 @@
 import pprint
 import struct
 import sys
+from typing import List
 
 import numpy as np
 
@@ -203,7 +204,7 @@ def resolve(self, key, is_local):
 
                 raise UndefinedVariableError(key, is_local)
 
-    def swapkey(self, old_key, new_key, new_value=None):
+    def swapkey(self, old_key: str, new_key: str, new_value=None):
         """
         Replace a variable name, with a potentially new value.
 
@@ -228,7 +229,7 @@ def swapkey(self, old_key, new_key, new_value=None):
                 mapping[new_key] = new_value
                 return
 
-    def _get_vars(self, stack, scopes):
+    def _get_vars(self, stack, scopes: List[str]):
         """
         Get specifically scoped variables from a list of stack frames.
 

pandas/io/pytables.py

@@ -2040,6 +2040,7 @@ def convert(
             Table row number: the end of the sub-selection. Values larger than
             the underlying table's row count are normalized to that.
         """
+        assert self.table is not None  # for mypy
 
         assert self.table is not None
 
@@ -3417,7 +3418,7 @@ def data_orientation(self):
             )
         )
 
-    def queryables(self):
+    def queryables(self) -> Dict[str, Any]:
         """ return a dict of the kinds allowable columns for this object """
 
         # compute the values_axes queryables