TYP: core.computation (#41007)

Author: jbrockmendel

pandas/core/computation/engines.py

@@ -12,6 +12,7 @@
     align_terms,
     reconstruct_object,
 )
+from pandas.core.computation.expr import Expr
 from pandas.core.computation.ops import (
     MATHOPS,
     REDUCTIONS,
@@ -26,13 +27,13 @@ class NumExprClobberingError(NameError):
     pass
 
 
-def _check_ne_builtin_clash(expr):
+def _check_ne_builtin_clash(expr: Expr) -> None:
     """
     Attempt to prevent foot-shooting in a helpful way.
 
     Parameters
     ----------
-    terms : Term
+    expr : Expr
         Terms can contain
     """
     names = expr.names

pandas/core/computation/eval.py

@@ -1,9 +1,9 @@
 """
 Top level ``eval`` module.
 """
+from __future__ import annotations
 
 import tokenize
-from typing import Optional
 import warnings
 
 from pandas._libs.lib import no_default
@@ -14,13 +14,14 @@
     PARSERS,
     Expr,
 )
+from pandas.core.computation.ops import BinOp
 from pandas.core.computation.parsing import tokenize_string
 from pandas.core.computation.scope import ensure_scope
 
 from pandas.io.formats.printing import pprint_thing
 
 
-def _check_engine(engine: Optional[str]) -> str:
+def _check_engine(engine: str | None) -> str:
     """
     Make sure a valid engine is passed.
 
@@ -161,9 +162,9 @@ def _check_for_locals(expr: str, stack_level: int, parser: str):
 
 
 def eval(
-    expr,
-    parser="pandas",
-    engine: Optional[str] = None,
+    expr: str | BinOp,  # we leave BinOp out of the docstr bc it isn't for users
+    parser: str = "pandas",
+    engine: str | None = None,
     truediv=no_default,
     local_dict=None,
     global_dict=None,
@@ -309,10 +310,12 @@ def eval(
             stacklevel=2,
         )
 
+    exprs: list[str | BinOp]
     if isinstance(expr, str):
         _check_expression(expr)
         exprs = [e.strip() for e in expr.splitlines() if e.strip() != ""]
     else:
+        # ops.BinOp; for internal compat, not intended to be passed by users
         exprs = [expr]
     multi_line = len(exprs) > 1
 

pandas/core/computation/pytables.py

@@ -546,6 +546,7 @@ class PyTablesExpr(expr.Expr):
 
     _visitor: PyTablesExprVisitor | None
     env: PyTablesScope
+    expr: str
 
     def __init__(
         self,
@@ -570,7 +571,7 @@ def __init__(
             local_dict = where.env.scope
             _where = where.expr
 
-        elif isinstance(where, (list, tuple)):
+        elif is_list_like(where):
             where = list(where)
             for idx, w in enumerate(where):
                 if isinstance(w, PyTablesExpr):
@@ -580,6 +581,7 @@ def __init__(
                     where[idx] = w
             _where = " & ".join(f"({w})" for w in com.flatten(where))
         else:
+            # _validate_where ensures we otherwise have a string
             _where = where
 
         self.expr = _where

pandas/core/computation/scope.py

@@ -106,9 +106,13 @@ class Scope:
     """
 
     __slots__ = ["level", "scope", "target", "resolvers", "temps"]
+    level: int
+    scope: DeepChainMap
+    resolvers: DeepChainMap
+    temps: dict
 
     def __init__(
-        self, level, global_dict=None, local_dict=None, resolvers=(), target=None
+        self, level: int, global_dict=None, local_dict=None, resolvers=(), target=None
     ):
         self.level = level + 1
 
@@ -146,8 +150,7 @@ def __init__(
 
         # assumes that resolvers are going from outermost scope to inner
         if isinstance(local_dict, Scope):
-            # error: Cannot determine type of 'resolvers'
-            resolvers += tuple(local_dict.resolvers.maps)  # type: ignore[has-type]
+            resolvers += tuple(local_dict.resolvers.maps)
         self.resolvers = DeepChainMap(*resolvers)
         self.temps = {}
 
@@ -212,7 +215,7 @@ def resolve(self, key: str, is_local: bool):
 
                 raise UndefinedVariableError(key, is_local) from err
 
-    def swapkey(self, old_key: str, new_key: str, new_value=None):
+    def swapkey(self, old_key: str, new_key: str, new_value=None) -> None:
         """
         Replace a variable name, with a potentially new value.
 
@@ -238,7 +241,7 @@ def swapkey(self, old_key: str, new_key: str, new_value=None):
                 mapping[new_key] = new_value  # type: ignore[index]
                 return
 
-    def _get_vars(self, stack, scopes: list[str]):
+    def _get_vars(self, stack, scopes: list[str]) -> None:
         """
         Get specifically scoped variables from a list of stack frames.
 
@@ -263,7 +266,7 @@ def _get_vars(self, stack, scopes: list[str]):
                 # scope after the loop
                 del frame
 
-    def _update(self, level: int):
+    def _update(self, level: int) -> None:
         """
         Update the current scope by going back `level` levels.
 
@@ -313,7 +316,7 @@ def ntemps(self) -> int:
         return len(self.temps)
 
     @property
-    def full_scope(self):
+    def full_scope(self) -> DeepChainMap:
         """
         Return the full scope for use with passing to engines transparently
         as a mapping.

pandas/core/reshape/pivot.py

@@ -180,7 +180,6 @@ def __internal_pivot_table(
                     # TODO: why does test_pivot_table_doctest_case fail if
                     # we don't do this apparently-unnecessary setitem?
                     agged[v] = agged[v]
-                    pass
                 else:
                     agged[v] = maybe_downcast_to_dtype(agged[v], data[v].dtype)
 

pandas/tests/io/pytables/test_store.py

@@ -673,17 +673,20 @@ def test_coordinates(setup_path):
         tm.assert_frame_equal(result, expected)
 
         # invalid
-        msg = "cannot process expression"
-        with pytest.raises(ValueError, match=msg):
+        msg = (
+            "where must be passed as a string, PyTablesExpr, "
+            "or list-like of PyTablesExpr"
+        )
+        with pytest.raises(TypeError, match=msg):
             store.select("df", where=np.arange(len(df), dtype="float64"))
 
-        with pytest.raises(ValueError, match=msg):
+        with pytest.raises(TypeError, match=msg):
             store.select("df", where=np.arange(len(df) + 1))
 
-        with pytest.raises(ValueError, match=msg):
+        with pytest.raises(TypeError, match=msg):
             store.select("df", where=np.arange(len(df)), start=5)
 
-        with pytest.raises(ValueError, match=msg):
+        with pytest.raises(TypeError, match=msg):
             store.select("df", where=np.arange(len(df)), start=5, stop=10)
 
         # selection with filter