TYP: misc annotations (pandas-dev#56667)

Author: twoertwein

pandas/_testing/_warnings.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from contextlib import (
+    AbstractContextManager,
     contextmanager,
     nullcontext,
 )
@@ -112,7 +113,9 @@ class for all warnings. To raise multiple types of exceptions,
                 )
 
 
-def maybe_produces_warning(warning: type[Warning], condition: bool, **kwargs):
+def maybe_produces_warning(
+    warning: type[Warning], condition: bool, **kwargs
+) -> AbstractContextManager:
     """
     Return a context manager that possibly checks a warning based on the condition
     """

pandas/compat/_optional.py

@@ -2,7 +2,11 @@
 
 import importlib
 import sys
-from typing import TYPE_CHECKING
+from typing import (
+    TYPE_CHECKING,
+    Literal,
+    overload,
+)
 import warnings
 
 from pandas.util._exceptions import find_stack_level
@@ -82,12 +86,35 @@ def get_version(module: types.ModuleType) -> str:
     return version
 
 
+@overload
+def import_optional_dependency(
+    name: str,
+    extra: str = ...,
+    min_version: str | None = ...,
+    *,
+    errors: Literal["raise"] = ...,
+) -> types.ModuleType:
+    ...
+
+
+@overload
+def import_optional_dependency(
+    name: str,
+    extra: str = ...,
+    min_version: str | None = ...,
+    *,
+    errors: Literal["warn", "ignore"],
+) -> types.ModuleType | None:
+    ...
+
+
 def import_optional_dependency(
     name: str,
     extra: str = "",
-    errors: str = "raise",
     min_version: str | None = None,
-):
+    *,
+    errors: Literal["raise", "warn", "ignore"] = "raise",
+) -> types.ModuleType | None:
     """
     Import an optional dependency.
 

pandas/core/arrays/_arrow_string_mixins.py

@@ -1,6 +1,9 @@
 from __future__ import annotations
 
-from typing import Literal
+from typing import (
+    TYPE_CHECKING,
+    Literal,
+)
 
 import numpy as np
 
@@ -10,6 +13,9 @@
     import pyarrow as pa
     import pyarrow.compute as pc
 
+if TYPE_CHECKING:
+    from pandas._typing import Self
+
 
 class ArrowStringArrayMixin:
     _pa_array = None
@@ -22,7 +28,7 @@ def _str_pad(
         width: int,
         side: Literal["left", "right", "both"] = "left",
         fillchar: str = " ",
-    ):
+    ) -> Self:
         if side == "left":
             pa_pad = pc.utf8_lpad
         elif side == "right":
@@ -35,7 +41,7 @@ def _str_pad(
             )
         return type(self)(pa_pad(self._pa_array, width=width, padding=fillchar))
 
-    def _str_get(self, i: int):
+    def _str_get(self, i: int) -> Self:
         lengths = pc.utf8_length(self._pa_array)
         if i >= 0:
             out_of_bounds = pc.greater_equal(i, lengths)
@@ -59,7 +65,7 @@ def _str_get(self, i: int):
 
     def _str_slice_replace(
         self, start: int | None = None, stop: int | None = None, repl: str | None = None
-    ):
+    ) -> Self:
         if repl is None:
             repl = ""
         if start is None:
@@ -68,13 +74,13 @@ def _str_slice_replace(
             stop = np.iinfo(np.int64).max
         return type(self)(pc.utf8_replace_slice(self._pa_array, start, stop, repl))
 
-    def _str_capitalize(self):
+    def _str_capitalize(self) -> Self:
         return type(self)(pc.utf8_capitalize(self._pa_array))
 
-    def _str_title(self):
+    def _str_title(self) -> Self:
         return type(self)(pc.utf8_title(self._pa_array))
 
-    def _str_swapcase(self):
+    def _str_swapcase(self) -> Self:
         return type(self)(pc.utf8_swapcase(self._pa_array))
 
     def _str_removesuffix(self, suffix: str):

pandas/core/dtypes/inference.py

@@ -36,7 +36,7 @@
 is_iterator = lib.is_iterator
 
 
-def is_number(obj) -> TypeGuard[Number | np.number]:
+def is_number(obj: object) -> TypeGuard[Number | np.number]:
     """
     Check if the object is a number.
 
@@ -77,7 +77,7 @@ def is_number(obj) -> TypeGuard[Number | np.number]:
     return isinstance(obj, (Number, np.number))
 
 
-def iterable_not_string(obj) -> bool:
+def iterable_not_string(obj: object) -> bool:
     """
     Check if the object is an iterable but not a string.
 
@@ -102,7 +102,7 @@ def iterable_not_string(obj) -> bool:
     return isinstance(obj, abc.Iterable) and not isinstance(obj, str)
 
 
-def is_file_like(obj) -> bool:
+def is_file_like(obj: object) -> bool:
     """
     Check if the object is a file-like object.
 
@@ -138,7 +138,7 @@ def is_file_like(obj) -> bool:
     return bool(hasattr(obj, "__iter__"))
 
 
-def is_re(obj) -> TypeGuard[Pattern]:
+def is_re(obj: object) -> TypeGuard[Pattern]:
     """
     Check if the object is a regex pattern instance.
 
@@ -163,7 +163,7 @@ def is_re(obj) -> TypeGuard[Pattern]:
     return isinstance(obj, Pattern)
 
 
-def is_re_compilable(obj) -> bool:
+def is_re_compilable(obj: object) -> bool:
     """
     Check if the object can be compiled into a regex pattern instance.
 
@@ -185,14 +185,14 @@ def is_re_compilable(obj) -> bool:
     False
     """
     try:
-        re.compile(obj)
+        re.compile(obj)  # type: ignore[call-overload]
     except TypeError:
         return False
     else:
         return True
 
 
-def is_array_like(obj) -> bool:
+def is_array_like(obj: object) -> bool:
     """
     Check if the object is array-like.
 
@@ -224,7 +224,7 @@ def is_array_like(obj) -> bool:
     return is_list_like(obj) and hasattr(obj, "dtype")
 
 
-def is_nested_list_like(obj) -> bool:
+def is_nested_list_like(obj: object) -> bool:
     """
     Check if the object is list-like, and that all of its elements
     are also list-like.
@@ -265,12 +265,13 @@ def is_nested_list_like(obj) -> bool:
     return (
         is_list_like(obj)
         and hasattr(obj, "__len__")
-        and len(obj) > 0
-        and all(is_list_like(item) for item in obj)
+        # need PEP 724 to handle these typing errors
+        and len(obj) > 0  # pyright: ignore[reportGeneralTypeIssues]
+        and all(is_list_like(item) for item in obj)  # type: ignore[attr-defined]
     )
 
 
-def is_dict_like(obj) -> bool:
+def is_dict_like(obj: object) -> bool:
     """
     Check if the object is dict-like.
 
@@ -303,7 +304,7 @@ def is_dict_like(obj) -> bool:
     )
 
 
-def is_named_tuple(obj) -> bool:
+def is_named_tuple(obj: object) -> bool:
     """
     Check if the object is a named tuple.
 
@@ -331,7 +332,7 @@ def is_named_tuple(obj) -> bool:
     return isinstance(obj, abc.Sequence) and hasattr(obj, "_fields")
 
 
-def is_hashable(obj) -> TypeGuard[Hashable]:
+def is_hashable(obj: object) -> TypeGuard[Hashable]:
     """
     Return True if hash(obj) will succeed, False otherwise.
 
@@ -370,7 +371,7 @@ def is_hashable(obj) -> TypeGuard[Hashable]:
         return True
 
 
-def is_sequence(obj) -> bool:
+def is_sequence(obj: object) -> bool:
     """
     Check if the object is a sequence of objects.
     String types are not included as sequences here.
@@ -394,14 +395,16 @@ def is_sequence(obj) -> bool:
     False
     """
     try:
-        iter(obj)  # Can iterate over it.
-        len(obj)  # Has a length associated with it.
+        # Can iterate over it.
+        iter(obj)  # type: ignore[call-overload]
+        # Has a length associated with it.
+        len(obj)  # type: ignore[arg-type]
         return not isinstance(obj, (str, bytes))
     except (TypeError, AttributeError):
         return False
 
 
-def is_dataclass(item) -> bool:
+def is_dataclass(item: object) -> bool:
     """
     Checks if the object is a data-class instance
 

pandas/core/flags.py

@@ -111,7 +111,7 @@ def __setitem__(self, key: str, value) -> None:
     def __repr__(self) -> str:
         return f"<Flags(allows_duplicate_labels={self.allows_duplicate_labels})>"
 
-    def __eq__(self, other) -> bool:
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, type(self)):
             return self.allows_duplicate_labels == other.allows_duplicate_labels
         return False

pandas/core/ops/invalid.py

@@ -4,7 +4,11 @@
 from __future__ import annotations
 
 import operator
-from typing import TYPE_CHECKING
+from typing import (
+    TYPE_CHECKING,
+    Callable,
+    NoReturn,
+)
 
 import numpy as np
 
@@ -41,7 +45,7 @@ def invalid_comparison(left, right, op) -> npt.NDArray[np.bool_]:
     return res_values
 
 
-def make_invalid_op(name: str):
+def make_invalid_op(name: str) -> Callable[..., NoReturn]:
     """
     Return a binary method that always raises a TypeError.
 
@@ -54,7 +58,7 @@ def make_invalid_op(name: str):
     invalid_op : function
     """
 
-    def invalid_op(self, other=None):
+    def invalid_op(self, other=None) -> NoReturn:
         typ = type(self).__name__
         raise TypeError(f"cannot perform {name} with this index type: {typ}")
 

pandas/core/ops/missing.py

@@ -30,7 +30,7 @@
 from pandas.core import roperator
 
 
-def _fill_zeros(result: np.ndarray, x, y):
+def _fill_zeros(result: np.ndarray, x, y) -> np.ndarray:
     """
     If this is a reversed op, then flip x,y
 

pandas/io/orc.py

@@ -2,7 +2,6 @@
 from __future__ import annotations
 
 import io
-from types import ModuleType
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -218,7 +217,7 @@ def to_orc(
 
     if engine != "pyarrow":
         raise ValueError("engine must be 'pyarrow'")
-    engine = import_optional_dependency(engine, min_version="10.0.1")
+    pyarrow = import_optional_dependency(engine, min_version="10.0.1")
     pa = import_optional_dependency("pyarrow")
     orc = import_optional_dependency("pyarrow.orc")
 
@@ -227,10 +226,9 @@ def to_orc(
         path = io.BytesIO()
     assert path is not None  # For mypy
     with get_handle(path, "wb", is_text=False) as handles:
-        assert isinstance(engine, ModuleType)  # For mypy
         try:
             orc.write_table(
-                engine.Table.from_pandas(df, preserve_index=index),
+                pyarrow.Table.from_pandas(df, preserve_index=index),
                 handles.handle,
                 **engine_kwargs,
             )

pandas/util/__init__.py

@@ -25,5 +25,5 @@ def __getattr__(key: str):
     raise AttributeError(f"module 'pandas.util' has no attribute '{key}'")
 
 
-def capitalize_first_letter(s):
+def capitalize_first_letter(s: str) -> str:
     return s[:1].upper() + s[1:]

pandas/util/_decorators.py

@@ -340,7 +340,7 @@ def wrapper(*args, **kwargs):
     return decorate
 
 
-def doc(*docstrings: None | str | Callable, **params) -> Callable[[F], F]:
+def doc(*docstrings: None | str | Callable, **params: object) -> Callable[[F], F]:
     """
     A decorator to take docstring templates, concatenate them and perform string
     substitution on them.

pyproject.toml

@@ -583,7 +583,6 @@ module = [
   "pandas._testing.*", # TODO
   "pandas.arrays", # TODO
   "pandas.compat.numpy.function", # TODO
-  "pandas.compat._optional", # TODO
   "pandas.compat.compressors", # TODO
   "pandas.compat.pickle_compat", # TODO
   "pandas.core._numba.executor", # TODO
@@ -602,7 +601,6 @@ module = [
   "pandas.core.dtypes.concat", # TODO
   "pandas.core.dtypes.dtypes", # TODO
   "pandas.core.dtypes.generic", # TODO
-  "pandas.core.dtypes.inference", # TODO
   "pandas.core.dtypes.missing", # TODO
   "pandas.core.groupby.categorical", # TODO
   "pandas.core.groupby.generic", # TODO