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

pandas/_typing.py

@@ -75,3 +75,5 @@
 
 # to maintain type information across generic functions and parametrization
 T = TypeVar("T")
+FuncType = Callable[..., Any]
+F = TypeVar("F", bound=FuncType)

pandas/compat/__init__.py

@@ -12,6 +12,8 @@
 import sys
 import warnings
 
+from pandas._typing import F
+
 PY37 = sys.version_info >= (3, 7)
 PY38 = sys.version_info >= (3, 8)
 PYPY = platform.python_implementation() == "PyPy"
@@ -25,7 +27,7 @@
 # found at https://bitbucket.org/gutworth/six
 
 
-def set_function_name(f, name, cls):
+def set_function_name(f: F, name: str, cls) -> F:
     """
     Bind the name/qualname attributes of the function.
     """

pandas/compat/chainmap.py

@@ -1,4 +1,6 @@
-from typing import ChainMap, MutableMapping, TypeVar, cast
+from typing import ChainMap, List, MutableMapping, TypeVar
+
+from pandas._typing import T
 
 _KT = TypeVar("_KT")
 _VT = TypeVar("_VT")
@@ -11,13 +13,17 @@ class DeepChainMap(ChainMap[_KT, _VT]):
     Only works when all passed mapping are mutable.
     """
 
+    # error: Incompatible types in assignment (expression has type
+    #  "List[MutableMapping[_KT, _VT]]", base class "ChainMap" defined the type
+    #  as "List[Mapping[_KT, _VT]]")  [assignment]
+    maps: List[MutableMapping[_KT, _VT]]  # type: ignore
+
     def __setitem__(self, key: _KT, value: _VT) -> None:
         for mapping in self.maps:
-            mutable_mapping = cast(MutableMapping[_KT, _VT], mapping)
-            if key in mutable_mapping:
-                mutable_mapping[key] = value
+            if key in mapping:
+                mapping[key] = value
                 return
-        cast(MutableMapping[_KT, _VT], self.maps[0])[key] = value
+        self.maps[0][key] = value
 
     def __delitem__(self, key: _KT) -> None:
         """
@@ -27,8 +33,11 @@ def __delitem__(self, key: _KT) -> None:
             If `key` doesn't exist.
         """
         for mapping in self.maps:
-            mutable_mapping = cast(MutableMapping[_KT, _VT], mapping)
             if key in mapping:
-                del mutable_mapping[key]
+                del mapping[key]
                 return
         raise KeyError(key)
+
+    # FIXME: return type of new_child incorrect in typeshed
+    def new_child(self: T, m) -> T:  # type: ignore
+        return super().new_child(m)  # type: ignore

pandas/core/algorithms.py

@@ -4,7 +4,7 @@
 """
 import operator
 from textwrap import dedent
-from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Callable, Dict, Optional, Tuple, Union
 from warnings import catch_warnings, simplefilter, warn
 
 import numpy as np
@@ -1074,6 +1074,9 @@ def __init__(self, obj, n: int, keep: str):
         if self.keep not in ("first", "last", "all"):
             raise ValueError('keep must be either "first", "last" or "all"')
 
+    def compute(self, method):
+        raise NotImplementedError
+
     def nlargest(self):
         return self.compute("nlargest")
 
@@ -1430,7 +1433,7 @@ def _take_nd_object(arr, indexer, out, axis: int, fill_value, mask_info):
 
 def _get_take_nd_function(
     ndim: int, arr_dtype, out_dtype, axis: int = 0, mask_info=None
-):
+) -> Callable:
     if ndim <= 2:
         tup = (arr_dtype.name, out_dtype.name)
         if ndim == 1:

pandas/core/arrays/interval.py

@@ -1,5 +1,6 @@
 from operator import le, lt
 import textwrap
+from typing import Any, Sequence, Tuple
 
 import numpy as np
 
@@ -432,30 +433,33 @@ def from_arrays(cls, left, right, closed="right", copy=False, dtype=None):
             ),
         )
     )
-    def from_tuples(cls, data, closed="right", copy=False, dtype=None):
+    def from_tuples(
+        cls, data: Sequence[Tuple[Any, Any]], closed="right", copy=False, dtype=None
+    ):
+        left: Sequence
+        right: Sequence
         if len(data):
             left, right = [], []
+            for d in data:
+                if isna(d):
+                    lhs = rhs = np.nan
+                else:
+                    name = cls.__name__
+                    try:
+                        # need list of length 2 tuples, e.g. [(0, 1), (1, 2), ...]
+                        lhs, rhs = d
+                    except ValueError:
+                        msg = f"{name}.from_tuples requires tuples of length 2, got {d}"
+                        raise ValueError(msg)
+                    except TypeError:
+                        msg = f"{name}.from_tuples received an invalid item, {d}"
+                        raise TypeError(msg)
+                left.append(lhs)
+                right.append(rhs)
         else:
             # ensure that empty data keeps input dtype
             left = right = data
 
-        for d in data:
-            if isna(d):
-                lhs = rhs = np.nan
-            else:
-                name = cls.__name__
-                try:
-                    # need list of length 2 tuples, e.g. [(0, 1), (1, 2), ...]
-                    lhs, rhs = d
-                except ValueError as err:
-                    msg = f"{name}.from_tuples requires tuples of length 2, got {d}"
-                    raise ValueError(msg) from err
-                except TypeError as err:
-                    msg = f"{name}.from_tuples received an invalid item, {d}"
-                    raise TypeError(msg) from err
-            left.append(lhs)
-            right.append(rhs)
-
         return cls.from_arrays(left, right, closed, copy=False, dtype=dtype)
 
     def _validate(self):

pandas/core/arrays/sparse/array.py

@@ -4,7 +4,7 @@
 from collections import abc
 import numbers
 import operator
-from typing import Any, Callable, Union
+from typing import Any, Callable, Type, TypeVar, Union
 import warnings
 
 import numpy as np
@@ -51,6 +51,8 @@
 
 import pandas.io.formats.printing as printing
 
+_SparseArrayT = TypeVar("_SparseArrayT", bound="SparseArray")
+
 # ----------------------------------------------------------------------------
 # Array
 
@@ -396,8 +398,11 @@ def __init__(
 
     @classmethod
     def _simple_new(
-        cls, sparse_array: np.ndarray, sparse_index: SparseIndex, dtype: SparseDtype
-    ) -> "SparseArray":
+        cls: Type[_SparseArrayT],
+        sparse_array: np.ndarray,
+        sparse_index: SparseIndex,
+        dtype: SparseDtype,
+    ) -> _SparseArrayT:
         new = object.__new__(cls)
         new._sparse_index = sparse_index
         new._sparse_values = sparse_array
@@ -821,11 +826,12 @@ def _get_val_at(self, loc):
             val = com.maybe_box_datetimelike(val, self.sp_values.dtype)
             return val
 
-    def take(self, indices, allow_fill=False, fill_value=None) -> "SparseArray":
+    def take(self, indices, allow_fill: bool = False, fill_value=None) -> "SparseArray":
         if is_scalar(indices):
             raise ValueError(f"'indices' must be an array, not a scalar '{indices}'.")
         indices = np.asarray(indices, dtype=np.int32)
 
+        result: Union[np.ndarray, "SparseArray"]
         if indices.size == 0:
             result = np.array([], dtype="object")
             kwargs = {"dtype": self.dtype}
@@ -1386,8 +1392,10 @@ def __abs__(self):
     # ------------------------------------------------------------------------
 
     @classmethod
-    def _create_unary_method(cls, op) -> Callable[["SparseArray"], "SparseArray"]:
-        def sparse_unary_method(self) -> "SparseArray":
+    def _create_unary_method(
+        cls: Type[_SparseArrayT], op
+    ) -> Callable[[_SparseArrayT], _SparseArrayT]:
+        def sparse_unary_method(self: _SparseArrayT) -> _SparseArrayT:
             fill_value = op(np.array(self.fill_value)).item()
             values = op(self.sp_values)
             dtype = SparseDtype(values.dtype, fill_value)
@@ -1479,14 +1487,14 @@ def cmp_method(self, other):
 
     @classmethod
     def _add_unary_ops(cls):
-        cls.__pos__ = cls._create_unary_method(operator.pos)
-        cls.__neg__ = cls._create_unary_method(operator.neg)
-        cls.__invert__ = cls._create_unary_method(operator.invert)
+        setattr(cls, "__pos__", cls._create_unary_method(operator.pos))
+        setattr(cls, "__neg__", cls._create_unary_method(operator.neg))
+        setattr(cls, "__invert__", cls._create_unary_method(operator.invert))
 
     @classmethod
     def _add_comparison_ops(cls):
-        cls.__and__ = cls._create_comparison_method(operator.and_)
-        cls.__or__ = cls._create_comparison_method(operator.or_)
+        setattr(cls, "__and__", cls._create_comparison_method(operator.and_))
+        setattr(cls, "__or__", cls._create_comparison_method(operator.or_))
         cls.__xor__ = cls._create_arithmetic_method(operator.xor)
         super()._add_comparison_ops()
 

pandas/core/base.py

@@ -4,7 +4,17 @@
 
 import builtins
 import textwrap
-from typing import Any, Dict, FrozenSet, List, Optional, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    FrozenSet,
+    List,
+    Optional,
+    Tuple,
+    Union,
+)
 
 import numpy as np
 
@@ -34,6 +44,9 @@
 from pandas.core.construction import create_series_with_explicit_dtype
 import pandas.core.nanops as nanops
 
+if TYPE_CHECKING:
+    from pandas import DataFrame, Series  # noqa: F401
+
 _shared_docs: Dict[str, str] = dict()
 _indexops_doc_kwargs = dict(
     klass="IndexOpsMixin",
@@ -244,7 +257,7 @@ def _gotitem(self, key, ndim: int, subset=None):
         """
         raise AbstractMethodError(self)
 
-    def aggregate(self, func, *args, **kwargs):
+    def aggregate(self, func, **kwargs):
         raise AbstractMethodError(self)
 
     agg = aggregate
@@ -279,13 +292,15 @@ def _try_aggregate_string_function(self, arg: str, *args, **kwargs):
             f"'{arg}' is not a valid function for '{type(self).__name__}' object"
         )
 
-    def _aggregate(self, arg, *args, **kwargs):
+    def _aggregate(
+        self, arg: Union[str, Dict], *args, **kwargs
+    ) -> Tuple[Optional[Union[Dict, "Series", "DataFrame"]], Optional[bool]]:
         """
         provide an implementation for the aggregators
 
         Parameters
         ----------
-        arg : string, dict, function
+        arg : str, dict, function
         *args : args to pass on to the function
         **kwargs : kwargs to pass on to the function
 
@@ -318,7 +333,7 @@ def _aggregate(self, arg, *args, **kwargs):
             # eg. {'A' : ['mean']}, normalize all to
             # be list-likes
             if any(is_aggregator(x) for x in arg.values()):
-                new_arg = {}
+                new_arg: Dict[Any, Union[Tuple, List, Dict]] = {}
                 for k, v in arg.items():
                     if not isinstance(v, (tuple, list, dict)):
                         new_arg[k] = [v]
@@ -374,7 +389,7 @@ def _agg_2dim(how):
                 colg = self._gotitem(self._selection, ndim=2, subset=obj)
                 return colg.aggregate(how)
 
-            def _agg(arg, func):
+            def _agg(arg: Dict, func: Callable) -> Dict:
                 """
                 run the aggregations over the arg with func
                 return a dict
@@ -386,7 +401,7 @@ def _agg(arg, func):
 
             # set the final keys
             keys = list(arg.keys())
-            result = {}
+            result: Any = {}
 
             if self._selection is not None:
 
@@ -460,7 +475,7 @@ def is_any_frame() -> bool:
                 return result, True
 
             # fall thru
-            from pandas import DataFrame, Series
+            from pandas import DataFrame, Series  # noqa: F811
 
             try:
                 result = DataFrame(result)
@@ -546,7 +561,7 @@ def _aggregate_multiple_funcs(self, arg, _axis):
             # we are concatting non-NDFrame objects,
             # e.g. a list of scalars
 
-            from pandas import Series
+            from pandas import Series  # noqa: F811
 
             result = Series(results, index=keys, name=self.name)
             if is_nested_object(result):

pandas/core/computation/expressions.py

@@ -6,6 +6,7 @@
 
 """
 
+from typing import Callable, List, Set
 import warnings
 
 import numpy as np
@@ -20,10 +21,10 @@
     import numexpr as ne
 
 _TEST_MODE = None
-_TEST_RESULT = None
+_TEST_RESULT: List[bool]
 _USE_NUMEXPR = _NUMEXPR_INSTALLED
-_evaluate = None
-_where = None
+_evaluate: Callable
+_where: Callable
 
 # the set of dtypes that we will allow pass to numexpr
 _ALLOWED_DTYPES = {
@@ -74,7 +75,7 @@ def _can_use_numexpr(op, op_str, a, b, dtype_check):
         # required min elements (otherwise we are adding overhead)
         if np.prod(a.shape) > _MIN_ELEMENTS:
             # check for dtype compatibility
-            dtypes = set()
+            dtypes: Set[str] = set()
             for o in [a, b]:
                 # Series implements dtypes, check for dimension count as well
                 if hasattr(o, "dtypes") and o.ndim > 1:
@@ -213,7 +214,7 @@ def where(cond, a, b, use_numexpr=True):
     return _where(cond, a, b) if use_numexpr else _where_standard(cond, a, b)
 
 
-def set_test_mode(v=True):
+def set_test_mode(v: bool = True) -> None:
     """
     Keeps track of whether numexpr was used.  Stores an additional ``True``
     for every successful use of evaluate with numexpr since the last