TYP: Update pyright (pandas-dev#56892)

* TYP: Update pyright

* isort

* int | float

Author: twoertwein

.pre-commit-config.yaml

@@ -132,7 +132,7 @@ repos:
         types: [python]
         stages: [manual]
         additional_dependencies: &pyright_dependencies
-        - [email protected].339
+        - [email protected].347
     -   id: pyright
         # note: assumes python env is setup and activated
         name: pyright reportGeneralTypeIssues

pandas/_config/config.py

@@ -88,7 +88,7 @@ class DeprecatedOption(NamedTuple):
 
 class RegisteredOption(NamedTuple):
     key: str
-    defval: object
+    defval: Any
     doc: str
     validator: Callable[[object], Any] | None
     cb: Callable[[str], Any] | None

pandas/_config/localization.py

@@ -10,7 +10,10 @@
 import platform
 import re
 import subprocess
-from typing import TYPE_CHECKING
+from typing import (
+    TYPE_CHECKING,
+    cast,
+)
 
 from pandas._config.config import options
 
@@ -152,7 +155,7 @@ def get_locales(
         out_locales = []
         for x in split_raw_locales:
             try:
-                out_locales.append(str(x, encoding=options.display.encoding))
+                out_locales.append(str(x, encoding=cast(str, options.display.encoding)))
             except UnicodeError:
                 # 'locale -a' is used to populated 'raw_locales' and on
                 # Redhat 7 Linux (and maybe others) prints locale names

pandas/_libs/lib.pyi

@@ -69,16 +69,26 @@ def fast_multiget(
     mapping: dict,
     keys: np.ndarray,  # object[:]
     default=...,
-) -> np.ndarray: ...
+) -> ArrayLike: ...
 def fast_unique_multiple_list_gen(gen: Generator, sort: bool = ...) -> list: ...
 def fast_unique_multiple_list(lists: list, sort: bool | None = ...) -> list: ...
+@overload
 def map_infer(
     arr: np.ndarray,
     f: Callable[[Any], Any],
-    convert: bool = ...,
+    *,
+    convert: Literal[False],
     ignore_na: bool = ...,
 ) -> np.ndarray: ...
 @overload
+def map_infer(
+    arr: np.ndarray,
+    f: Callable[[Any], Any],
+    *,
+    convert: bool = ...,
+    ignore_na: bool = ...,
+) -> ArrayLike: ...
+@overload
 def maybe_convert_objects(
     objects: npt.NDArray[np.object_],
     *,
@@ -164,14 +174,26 @@ def is_all_arraylike(obj: list) -> bool: ...
 # Functions which in reality take memoryviews
 
 def memory_usage_of_objects(arr: np.ndarray) -> int: ...  # object[:]  # np.int64
+@overload
 def map_infer_mask(
     arr: np.ndarray,
     f: Callable[[Any], Any],
     mask: np.ndarray,  # const uint8_t[:]
-    convert: bool = ...,
+    *,
+    convert: Literal[False],
     na_value: Any = ...,
     dtype: np.dtype = ...,
 ) -> np.ndarray: ...
+@overload
+def map_infer_mask(
+    arr: np.ndarray,
+    f: Callable[[Any], Any],
+    mask: np.ndarray,  # const uint8_t[:]
+    *,
+    convert: bool = ...,
+    na_value: Any = ...,
+    dtype: np.dtype = ...,
+) -> ArrayLike: ...
 def indices_fast(
     index: npt.NDArray[np.intp],
     labels: np.ndarray,  # const int64_t[:]

pandas/_libs/lib.pyx

@@ -2864,10 +2864,11 @@ def map_infer_mask(
         ndarray[object] arr,
         object f,
         const uint8_t[:] mask,
+        *,
         bint convert=True,
         object na_value=no_default,
         cnp.dtype dtype=np.dtype(object)
-) -> np.ndarray:
+) -> "ArrayLike":
     """
     Substitute for np.vectorize with pandas-friendly dtype inference.
 
@@ -2887,7 +2888,7 @@ def map_infer_mask(
 
     Returns
     -------
-    np.ndarray
+    np.ndarray or an ExtensionArray
     """
     cdef Py_ssize_t n = len(arr)
     result = np.empty(n, dtype=dtype)
@@ -2941,8 +2942,8 @@ def _map_infer_mask(
 @cython.boundscheck(False)
 @cython.wraparound(False)
 def map_infer(
-    ndarray arr, object f, bint convert=True, bint ignore_na=False
-) -> np.ndarray:
+    ndarray arr, object f, *, bint convert=True, bint ignore_na=False
+) -> "ArrayLike":
     """
     Substitute for np.vectorize with pandas-friendly dtype inference.
 
@@ -2956,7 +2957,7 @@ def map_infer(
 
     Returns
     -------
-    np.ndarray
+    np.ndarray or an ExtensionArray
     """
     cdef:
         Py_ssize_t i, n
@@ -3091,7 +3092,7 @@ def to_object_array_tuples(rows: object) -> np.ndarray:
 
 @cython.wraparound(False)
 @cython.boundscheck(False)
-def fast_multiget(dict mapping, object[:] keys, default=np.nan) -> np.ndarray:
+def fast_multiget(dict mapping, object[:] keys, default=np.nan) -> "ArrayLike":
     cdef:
         Py_ssize_t i, n = len(keys)
         object val

pandas/core/arrays/string_.py

@@ -4,6 +4,7 @@
     TYPE_CHECKING,
     ClassVar,
     Literal,
+    cast,
 )
 
 import numpy as np
@@ -637,7 +638,7 @@ def _str_map(
                 # error: Argument 1 to "dtype" has incompatible type
                 # "Union[ExtensionDtype, str, dtype[Any], Type[object]]"; expected
                 # "Type[object]"
-                dtype=np.dtype(dtype),  # type: ignore[arg-type]
+                dtype=np.dtype(cast(type, dtype)),
             )
 
             if not na_value_is_na:

pandas/core/arrays/string_arrow.py

@@ -7,6 +7,7 @@
     TYPE_CHECKING,
     Callable,
     Union,
+    cast,
 )
 import warnings
 
@@ -327,7 +328,7 @@ def _str_map(
                 # error: Argument 1 to "dtype" has incompatible type
                 # "Union[ExtensionDtype, str, dtype[Any], Type[object]]"; expected
                 # "Type[object]"
-                dtype=np.dtype(dtype),  # type: ignore[arg-type]
+                dtype=np.dtype(cast(type, dtype)),
             )
 
             if not na_value_is_na:
@@ -640,7 +641,7 @@ def _str_map(
                     mask.view("uint8"),
                     convert=False,
                     na_value=na_value,
-                    dtype=np.dtype(dtype),  # type: ignore[arg-type]
+                    dtype=np.dtype(cast(type, dtype)),
                 )
                 return result
 

pandas/core/computation/eval.py

@@ -4,7 +4,10 @@
 from __future__ import annotations
 
 import tokenize
-from typing import TYPE_CHECKING
+from typing import (
+    TYPE_CHECKING,
+    Any,
+)
 import warnings
 
 from pandas.util._exceptions import find_stack_level
@@ -177,7 +180,7 @@ def eval(
     level: int = 0,
     target=None,
     inplace: bool = False,
-):
+) -> Any:
     """
     Evaluate a Python expression as a string using various backends.
 

pandas/core/dtypes/missing.py

@@ -258,7 +258,9 @@ def _use_inf_as_na(key) -> None:
         globals()["INF_AS_NA"] = False
 
 
-def _isna_array(values: ArrayLike, inf_as_na: bool = False):
+def _isna_array(
+    values: ArrayLike, inf_as_na: bool = False
+) -> npt.NDArray[np.bool_] | NDFrame:
     """
     Return an array indicating which values of the input array are NaN / NA.
 
@@ -275,6 +277,7 @@ def _isna_array(values: ArrayLike, inf_as_na: bool = False):
         Array of boolean values denoting the NA status of each element.
     """
     dtype = values.dtype
+    result: npt.NDArray[np.bool_] | NDFrame
 
     if not isinstance(values, np.ndarray):
         # i.e. ExtensionArray

pandas/core/frame.py

@@ -9802,7 +9802,9 @@ def explode(
 
         return result.__finalize__(self, method="explode")
 
-    def unstack(self, level: IndexLabel = -1, fill_value=None, sort: bool = True):
+    def unstack(
+        self, level: IndexLabel = -1, fill_value=None, sort: bool = True
+    ) -> DataFrame | Series:
         """
         Pivot a level of the (necessarily hierarchical) index labels.
 

pandas/core/indexes/range.py

@@ -491,7 +491,7 @@ def copy(self, name: Hashable | None = None, deep: bool = False) -> Self:
         new_index = self._rename(name=name)
         return new_index
 
-    def _minmax(self, meth: str):
+    def _minmax(self, meth: str) -> int | float:
         no_steps = len(self) - 1
         if no_steps == -1:
             return np.nan
@@ -500,13 +500,13 @@ def _minmax(self, meth: str):
 
         return self.start + self.step * no_steps
 
-    def min(self, axis=None, skipna: bool = True, *args, **kwargs) -> int:
+    def min(self, axis=None, skipna: bool = True, *args, **kwargs) -> int | float:
         """The minimum value of the RangeIndex"""
         nv.validate_minmax_axis(axis)
         nv.validate_min(args, kwargs)
         return self._minmax("min")
 
-    def max(self, axis=None, skipna: bool = True, *args, **kwargs) -> int:
+    def max(self, axis=None, skipna: bool = True, *args, **kwargs) -> int | float:
         """The maximum value of the RangeIndex"""
         nv.validate_minmax_axis(axis)
         nv.validate_max(args, kwargs)

pandas/core/ops/common.py

@@ -40,7 +40,7 @@ def wrapper(method: F) -> F:
     return wrapper
 
 
-def _unpack_zerodim_and_defer(method, name: str):
+def _unpack_zerodim_and_defer(method: F, name: str) -> F:
     """
     Boilerplate for pandas conventions in arithmetic and comparison methods.
 
@@ -75,7 +75,9 @@ def new_method(self, other):
 
         return method(self, other)
 
-    return new_method
+    # error: Incompatible return value type (got "Callable[[Any, Any], Any]",
+    # expected "F")
+    return new_method  # type: ignore[return-value]
 
 
 def get_op_result_name(left, right):

pandas/core/reshape/pivot.py

@@ -568,7 +568,8 @@ def pivot(
     # error: Argument 1 to "unstack" of "DataFrame" has incompatible type "Union
     # [List[Any], ExtensionArray, ndarray[Any, Any], Index, Series]"; expected
     # "Hashable"
-    result = indexed.unstack(columns_listlike)  # type: ignore[arg-type]
+    # unstack with a MultiIndex returns a DataFrame
+    result = cast("DataFrame", indexed.unstack(columns_listlike))  # type: ignore[arg-type]
     result.index.names = [
         name if name is not lib.no_default else None for name in result.index.names
     ]