TYP: rename (TypeVar approach)

Author: twoertwein

pandas/_typing.py

@@ -141,7 +141,10 @@
 ]
 
 # For functions like rename that convert one label to another
-Renamer = Union[Mapping[Hashable, Any], Callable[[Hashable], Hashable]]
+HashableT = TypeVar("HashableT", bound=Hashable)
+HashableTa = TypeVar("HashableTa", bound=Hashable)
+HashableTb = TypeVar("HashableTb", bound=Hashable)
+Renamer = Union[Mapping[HashableT, Hashable], Callable[[HashableT], Hashable]]
 
 # to maintain type information across generic functions and parametrization
 T = TypeVar("T")

pandas/core/frame.py

@@ -58,6 +58,9 @@
     FloatFormatType,
     FormattersType,
     Frequency,
+    HashableT,
+    HashableTa,
+    HashableTb,
     IgnoreRaise,
     IndexKeyFunc,
     IndexLabel,
@@ -5030,17 +5033,62 @@ def drop(  # type: ignore[override]
             errors=errors,
         )
 
+    @overload
+    def rename(
+        self,
+        mapper: Renamer[HashableT] | None = ...,
+        *,
+        index: Renamer[HashableTa] | None = ...,
+        columns: Renamer[HashableTb] | None = ...,
+        axis: Axis | None = ...,
+        copy: bool = ...,
+        inplace: Literal[True],
+        level: Level | None = ...,
+        errors: IgnoreRaise = ...,
+    ) -> None:
+        ...
+
+    @overload
+    def rename(
+        self,
+        mapper: Renamer[HashableT] | None = ...,
+        *,
+        index: Renamer[HashableTa] | None = ...,
+        columns: Renamer[HashableTb] | None = ...,
+        axis: Axis | None = ...,
+        copy: bool = ...,
+        inplace: Literal[False] = ...,
+        level: Level | None = ...,
+        errors: IgnoreRaise = ...,
+    ) -> DataFrame:
+        ...
+
+    @overload
     def rename(
         self,
-        mapper: Renamer | None = None,
+        mapper: Renamer[HashableT] | None = ...,
         *,
-        index: Renamer | None = None,
-        columns: Renamer | None = None,
+        index: Renamer[HashableTa] | None = ...,
+        columns: Renamer[HashableTb] | None = ...,
+        axis: Axis | None = ...,
+        copy: bool = ...,
+        inplace: bool = ...,
+        level: Level | None = ...,
+        errors: IgnoreRaise = ...,
+    ) -> DataFrame | None:
+        ...
+
+    def rename(
+        self,
+        mapper: Renamer[HashableT] | None = None,
+        *,
+        index: Renamer[HashableTa] | None = None,
+        columns: Renamer[HashableTb] | None = None,
         axis: Axis | None = None,
         copy: bool = True,
         inplace: bool = False,
         level: Level | None = None,
-        errors: str = "ignore",
+        errors: IgnoreRaise = "ignore",
     ) -> DataFrame | None:
         """
         Alter axes labels.

pandas/core/generic.py

@@ -44,6 +44,9 @@
     DtypeArg,
     DtypeObj,
     FilePath,
+    HashableT,
+    HashableTa,
+    HashableTb,
     IgnoreRaise,
     IndexKeyFunc,
     IndexLabel,
@@ -974,10 +977,10 @@ def squeeze(self, axis=None):
 
     def _rename(
         self: NDFrameT,
-        mapper: Renamer | None = None,
+        mapper: Renamer[HashableT] | None = None,
         *,
-        index: Renamer | None = None,
-        columns: Renamer | None = None,
+        index: Renamer[HashableTa] | None = None,
+        columns: Renamer[HashableTb] | None = None,
         axis: Axis | None = None,
         copy: bool_t = True,
         inplace: bool_t = False,
@@ -1110,9 +1113,19 @@ def _rename(
         else:
             # use the mapper argument
             if axis and self._get_axis_number(axis) == 1:
-                columns = mapper
+                # error: Incompatible types in assignment (expression has type
+                # "Optional[Union[Mapping[HashableT, Hashable], Callable[
+                # [HashableT], Hashable]]]", variable has type "Optional[Union[
+                # Mapping[HashableTb, Hashable], Callable[[HashableTb], Hashable
+                # ]]]")
+                columns = mapper  # type: ignore[assignment]
             else:
-                index = mapper
+                # error: Incompatible types in assignment (expression has type
+                # "Optional[Union[Mapping[HashableT, Hashable], Callable[[
+                # HashableT], Hashable]]]", variable has type "Optional[Union[
+                # Mapping[HashableTa, Hashable], Callable[[HashableTa], Hashable
+                # ]]]")
+                index = mapper  # type: ignore[assignment]
 
         self._check_inplace_and_allows_duplicate_labels(inplace)
         result = self if inplace else self.copy(deep=copy)

pandas/core/groupby/groupby.py

@@ -2204,8 +2204,7 @@ def size(self) -> DataFrame | Series:
             result = self._obj_1d_constructor(result)
 
         if not self.as_index:
-            # Item "None" of "Optional[Series]" has no attribute "reset_index"
-            result = result.rename("size").reset_index()  # type: ignore[union-attr]
+            result = result.rename("size").reset_index()
 
         return self._reindex_output(result, fill_value=0)
 

pandas/core/series.py

@@ -38,10 +38,12 @@
     Dtype,
     DtypeObj,
     FillnaOptions,
+    HashableT,
     IgnoreRaise,
     IndexKeyFunc,
     Level,
     NaPosition,
+    Renamer,
     SingleManager,
     SortKind,
     StorageOptions,
@@ -4617,15 +4619,54 @@ def align(
             broadcast_axis=broadcast_axis,
         )
 
+    @overload
     def rename(
         self,
-        index=None,
+        index: Renamer[HashableT] | Hashable | None = ...,
         *,
-        axis=None,
-        copy=True,
-        inplace=False,
-        level=None,
-        errors="ignore",
+        axis: Axis | None = ...,
+        copy: bool = ...,
+        inplace: Literal[True],
+        level: Level | None = ...,
+        errors: IgnoreRaise = ...,
+    ) -> None:
+        ...
+
+    @overload
+    def rename(
+        self,
+        index: Renamer[HashableT] | Hashable | None = ...,
+        *,
+        axis: Axis | None = ...,
+        copy: bool = ...,
+        inplace: Literal[False] = ...,
+        level: Level | None = ...,
+        errors: IgnoreRaise = ...,
+    ) -> Series:
+        ...
+
+    @overload
+    def rename(
+        self,
+        index: Renamer[HashableT] | Hashable | None = ...,
+        *,
+        axis: Axis | None = ...,
+        copy: bool = ...,
+        inplace: bool = ...,
+        level: Level | None = ...,
+        errors: IgnoreRaise = ...,
+    ) -> Series | None:
+        ...
+
+    def rename(
+        self,
+        index: Renamer[HashableT] | Hashable | None = None,
+        *,
+        axis: Axis | None = None,
+        copy: bool = True,
+        inplace: bool = False,
+        level: Level | None = None,
+        errors: IgnoreRaise = "ignore",
     ) -> Series | None:
         """
         Alter Series index labels or name.
@@ -4691,8 +4732,16 @@ def rename(
             axis = self._get_axis_number(axis)
 
         if callable(index) or is_dict_like(index):
+            # error: Argument 1 to "_rename" of "NDFrame" has incompatible type
+            # "Union[Union[Mapping[Hashable, Any], Callable[[Hashable], Hashable]],
+            # Hashable, None]"; expected "Union[Mapping[Hashable, Any],
+            # Callable[[Hashable], Hashable], None]"
             return super()._rename(
-                index, copy=copy, inplace=inplace, level=level, errors=errors
+                index,  # type: ignore[arg-type]
+                copy=copy,
+                inplace=inplace,
+                level=level,
+                errors=errors,
             )
         else:
             return self._set_name(index, inplace=inplace)

pandas/io/json/_normalize.py

@@ -9,7 +9,9 @@
 import copy
 from typing import (
     Any,
+    Callable,
     DefaultDict,
+    Hashable,
     Iterable,
 )
 
@@ -520,11 +522,8 @@ def _recursive_extract(data, path, seen_meta, level=0):
     result = DataFrame(records)
 
     if record_prefix is not None:
-        # Incompatible types in assignment (expression has type "Optional[DataFrame]",
-        # variable has type "DataFrame")
-        result = result.rename(  # type: ignore[assignment]
-            columns=lambda x: f"{record_prefix}{x}"
-        )
+        rename_fun: Callable[[Hashable], str] = lambda x: f"{record_prefix}{x}"
+        result = result.rename(columns=rename_fun)
 
     # Data types, a problem
     for k, v in meta_vals.items():