Add overloads for NDFrame.rename_axis. (#85)

* Add overloads for NDFrame.rename_axis. * Add test for Series.rename_axis
VirtusLab · Sep 17, 2021 · e7904db · e7904db
1 parent 2f79376
commit e7904db
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/tests/snippets/test_series.py b/tests/snippets/test_series.py
@@ -418,3 +418,7 @@ def test_types_eq() -> None:
     res1: pd.Series = s1 == 1
     s2 = pd.Series([1, 2, 4])
     res2: pd.Series = s1 == s2
+
+
+def test_types_rename_axis() -> None:
+    s: pd.Series = pd.Series([1, 2, 3]).rename_axis("A")
diff --git a/third_party/3/pandas/core/generic.pyi b/third_party/3/pandas/core/generic.pyi
@@ -76,7 +76,10 @@ class NDFrame(PandasObject, SelectionMixin, indexing.IndexingMixin):
     def rename(self, mapper: Optional[Renamer]=..., *, index: Optional[Renamer]=..., columns: Optional[Renamer]=..., axis: Optional[Axis]=..., copy: bool_t=..., inplace: Literal[False] = ..., level: Optional[Level]=..., errors: ErrorsStrategy=...) -> FrameOrSeries: ...
     @overload
     def rename(self, mapper: Optional[Renamer]=..., *, index: Optional[Renamer]=..., columns: Optional[Renamer]=..., axis: Optional[Axis]=..., copy: bool_t=..., inplace: Literal[True], level: Optional[Level]=..., errors: ErrorsStrategy=...) -> None: ...
-    def rename_axis(self, mapper: Any = ..., **kwargs: Any) -> Optional[FrameOrSeries]: ...
+    @overload
+    def rename_axis(self: Self, mapper: Any = ..., inplace: Literal[False] = ..., **kwargs: Any) -> Self: ...
+    @overload
+    def rename_axis(self, inplace: Literal[True], mapper: Any = ..., **kwargs: Any) -> None: ...
     def equals(self, other: Any) -> Any: ...
     @property
     def empty(self) -> bool_t: ...