TYP: misc annotations (#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.