TYP: core.computation mostly (#29652)

Author: jbrockmendel

pandas/_libs/internals.pyx

@@ -85,7 +85,7 @@ cdef class BlockPlacement:
             return iter(self._as_array)
 
     @property
-    def as_slice(self):
+    def as_slice(self) -> slice:
         cdef:
             slice s = self._ensure_has_slice()
         if s is None:
@@ -118,7 +118,7 @@ cdef class BlockPlacement:
         return self._as_array
 
     @property
-    def is_slice_like(self):
+    def is_slice_like(self) -> bool:
         cdef:
             slice s = self._ensure_has_slice()
         return s is not None
@@ -441,7 +441,7 @@ def get_blkno_indexers(int64_t[:] blknos, bint group=True):
                 yield blkno, result
 
 
-def get_blkno_placements(blknos, group=True):
+def get_blkno_placements(blknos, group: bool = True):
     """
 
     Parameters

pandas/_libs/sparse.pyx

@@ -57,7 +57,7 @@ cdef class IntIndex(SparseIndex):
         return output
 
     @property
-    def nbytes(self):
+    def nbytes(self) -> int:
         return self.indices.nbytes
 
     def check_integrity(self):
@@ -91,7 +91,7 @@ cdef class IntIndex(SparseIndex):
         if not monotonic:
             raise ValueError("Indices must be strictly increasing")
 
-    def equals(self, other):
+    def equals(self, other) -> bool:
         if not isinstance(other, IntIndex):
             return False
 
@@ -103,7 +103,7 @@ cdef class IntIndex(SparseIndex):
         return same_length and same_indices
 
     @property
-    def ngaps(self):
+    def ngaps(self) -> int:
         return self.length - self.npoints
 
     def to_int_index(self):
@@ -348,11 +348,11 @@ cdef class BlockIndex(SparseIndex):
         return output
 
     @property
-    def nbytes(self):
+    def nbytes(self) -> int:
         return self.blocs.nbytes + self.blengths.nbytes
 
     @property
-    def ngaps(self):
+    def ngaps(self) -> int:
         return self.length - self.npoints
 
     cpdef check_integrity(self):
@@ -388,7 +388,7 @@ cdef class BlockIndex(SparseIndex):
             if blengths[i] == 0:
                 raise ValueError(f'Zero-length block {i}')
 
-    def equals(self, other):
+    def equals(self, other) -> bool:
         if not isinstance(other, BlockIndex):
             return False
 

pandas/_libs/writers.pyx

@@ -70,7 +70,7 @@ def write_csv_rows(list data, ndarray data_index,
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
-def convert_json_to_lines(object arr):
+def convert_json_to_lines(arr: object) -> str:
     """
     replace comma separated json with line feeds, paying special attention
     to quotes & brackets

pandas/core/computation/align.py

@@ -33,7 +33,7 @@ def _zip_axes_from_type(typ, new_axes):
     return axes
 
 
-def _any_pandas_objects(terms):
+def _any_pandas_objects(terms) -> bool:
     """Check a sequence of terms for instances of PandasObject."""
     return any(isinstance(term.value, PandasObject) for term in terms)
 

pandas/core/computation/engines.py

@@ -46,8 +46,9 @@ def __init__(self, expr):
         self.aligned_axes = None
         self.result_type = None
 
-    def convert(self):
-        """Convert an expression for evaluation.
+    def convert(self) -> str:
+        """
+        Convert an expression for evaluation.
 
         Defaults to return the expression as a string.
         """
@@ -75,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
@@ -104,7 +105,7 @@ class NumExprEngine(AbstractEngine):
     def __init__(self, expr):
         super().__init__(expr)
 
-    def convert(self):
+    def convert(self) -> str:
         return str(super().convert())
 
     def _evaluate(self):

pandas/core/computation/ops.py

@@ -70,6 +70,7 @@ def __new__(cls, name, env, side=None, encoding=None):
         return supr_new(klass)
 
     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
         self.env = env
         self.side = side
@@ -79,7 +80,7 @@ def __init__(self, name, env, side=None, encoding=None):
         self.encoding = encoding
 
     @property
-    def local_name(self):
+    def local_name(self) -> str:
         return self.name.replace(_LOCAL_TAG, "")
 
     def __repr__(self) -> str:
@@ -120,7 +121,7 @@ def update(self, value):
         self.value = value
 
     @property
-    def is_scalar(self):
+    def is_scalar(self) -> bool:
         return is_scalar(self._value)
 
     @property
@@ -139,14 +140,14 @@ def type(self):
     return_type = type
 
     @property
-    def raw(self):
+    def raw(self) -> str:
         return pprint_thing(
             "{0}(name={1!r}, type={2})"
             "".format(self.__class__.__name__, self.name, self.type)
         )
 
     @property
-    def is_datetime(self):
+    def is_datetime(self) -> bool:
         try:
             t = self.type.type
         except AttributeError:
@@ -220,7 +221,7 @@ def return_type(self):
         return _result_type_many(*(term.type for term in com.flatten(self)))
 
     @property
-    def has_invalid_return_type(self):
+    def has_invalid_return_type(self) -> bool:
         types = self.operand_types
         obj_dtype_set = frozenset([np.dtype("object")])
         return self.return_type == object and types - obj_dtype_set
@@ -230,11 +231,11 @@ def operand_types(self):
         return frozenset(term.type for term in com.flatten(self))
 
     @property
-    def is_scalar(self):
+    def is_scalar(self) -> bool:
         return all(operand.is_scalar for operand in self.operands)
 
     @property
-    def is_datetime(self):
+    def is_datetime(self) -> bool:
         try:
             t = self.return_type.type
         except AttributeError:
@@ -339,7 +340,7 @@ def _cast_inplace(terms, acceptable_dtypes, dtype):
         term.update(new_value)
 
 
-def is_term(obj):
+def is_term(obj) -> bool:
     return isinstance(obj, Term)
 
 
@@ -354,7 +355,7 @@ class BinOp(Op):
     right : Term or Op
     """
 
-    def __init__(self, op, lhs, rhs, **kwargs):
+    def __init__(self, op: str, lhs, rhs, **kwargs):
         super().__init__(op, (lhs, rhs))
         self.lhs = lhs
         self.rhs = rhs
@@ -396,7 +397,7 @@ def __call__(self, env):
 
         return self.func(left, right)
 
-    def evaluate(self, env, engine, parser, term_type, eval_in_python):
+    def evaluate(self, env, engine: str, parser, term_type, eval_in_python):
         """
         Evaluate a binary operation *before* being passed to the engine.
 
@@ -488,7 +489,7 @@ def _disallow_scalar_only_bool_ops(self):
             raise NotImplementedError("cannot evaluate scalar only bool ops")
 
 
-def isnumeric(dtype):
+def isnumeric(dtype) -> bool:
     return issubclass(np.dtype(dtype).type, np.number)
 
 
@@ -505,8 +506,8 @@ class Div(BinOp):
         regardless of the value of ``truediv``.
     """
 
-    def __init__(self, lhs, rhs, truediv, *args, **kwargs):
-        super().__init__("/", lhs, rhs, *args, **kwargs)
+    def __init__(self, lhs, rhs, truediv: bool, **kwargs):
+        super().__init__("/", lhs, rhs, **kwargs)
 
         if not isnumeric(lhs.return_type) or not isnumeric(rhs.return_type):
             raise TypeError(
@@ -541,7 +542,7 @@ class UnaryOp(Op):
         * If no function associated with the passed operator token is found.
     """
 
-    def __init__(self, op, operand):
+    def __init__(self, op: str, operand):
         super().__init__(op, (operand,))
         self.operand = operand
 
@@ -561,7 +562,7 @@ def __repr__(self) -> str:
         return pprint_thing("{0}({1})".format(self.op, self.operand))
 
     @property
-    def return_type(self):
+    def return_type(self) -> np.dtype:
         operand = self.operand
         if operand.return_type == np.dtype("bool"):
             return np.dtype("bool")
@@ -588,7 +589,7 @@ def __repr__(self) -> str:
 
 
 class FuncNode:
-    def __init__(self, name):
+    def __init__(self, name: str):
         from pandas.core.computation.check import _NUMEXPR_INSTALLED, _NUMEXPR_VERSION
 
         if name not in _mathops or (

pandas/core/computation/pytables.py

@@ -2,6 +2,7 @@
 
 import ast
 from functools import partial
+from typing import Optional
 
 import numpy as np
 
@@ -129,12 +130,12 @@ def conform(self, rhs):
         return rhs
 
     @property
-    def is_valid(self):
+    def is_valid(self) -> bool:
         """ return True if this is a valid field """
         return self.lhs in self.queryables
 
     @property
-    def is_in_table(self):
+    def is_in_table(self) -> bool:
         """ return True if this is a valid column name for generation (e.g. an
         actual column in the table) """
         return self.queryables.get(self.lhs) is not None
@@ -154,12 +155,12 @@ def metadata(self):
         """ the metadata of my field """
         return getattr(self.queryables.get(self.lhs), "metadata", None)
 
-    def generate(self, v):
+    def generate(self, v) -> str:
         """ create and return the op string for this TermValue """
         val = v.tostring(self.encoding)
         return "({lhs} {op} {val})".format(lhs=self.lhs, op=self.op, val=val)
 
-    def convert_value(self, v):
+    def convert_value(self, v) -> "TermValue":
         """ convert the expression that is in the term to something that is
         accepted by pytables """
 
@@ -279,7 +280,7 @@ def evaluate(self):
 
         return self
 
-    def generate_filter_op(self, invert=False):
+    def generate_filter_op(self, invert: bool = False):
         if (self.op == "!=" and not invert) or (self.op == "==" and invert):
             return lambda axis, vals: ~axis.isin(vals)
         else:
@@ -505,7 +506,7 @@ class Expr(expr.Expr):
     "major_axis>=20130101"
     """
 
-    def __init__(self, where, queryables=None, encoding=None, scope_level=0):
+    def __init__(self, where, queryables=None, encoding=None, scope_level: int = 0):
 
         where = _validate_where(where)
 
@@ -520,18 +521,21 @@ def __init__(self, where, queryables=None, encoding=None, scope_level=0):
 
         if isinstance(where, Expr):
             local_dict = where.env.scope
-            where = where.expr
+            _where = where.expr
 
         elif isinstance(where, (list, tuple)):
+            where = list(where)
             for idx, w in enumerate(where):
                 if isinstance(w, Expr):
                     local_dict = w.env.scope
                 else:
                     w = _validate_where(w)
                     where[idx] = w
-            where = " & ".join(map("({})".format, com.flatten(where)))  # noqa
+            _where = " & ".join(map("({})".format, com.flatten(where)))
+        else:
+            _where = where
 
-        self.expr = where
+        self.expr = _where
         self.env = Scope(scope_level + 1, local_dict=local_dict)
 
         if queryables is not None and isinstance(self.expr, str):
@@ -574,7 +578,7 @@ def evaluate(self):
 class TermValue:
     """ hold a term value the we use to construct a condition/filter """
 
-    def __init__(self, value, converted, kind):
+    def __init__(self, value, converted, kind: Optional[str]):
         self.value = value
         self.converted = converted
         self.kind = kind
@@ -593,7 +597,7 @@ def tostring(self, encoding):
         return self.converted
 
 
-def maybe_expression(s):
+def maybe_expression(s) -> bool:
     """ loose checking if s is a pytables-acceptable expression """
     if not isinstance(s, str):
         return False