TYP: NDFrame.__round__ (#792)

* TYP: NDFrame.__round__

* more Self

* a few more Selfs

Author: twoertwein

.pre-commit-config.yaml

@@ -11,7 +11,7 @@ repos:
     hooks:
     -   id: isort
 -   repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.0.291
+    rev: v0.0.292
     hooks:
     -   id: ruff
         args: [
@@ -23,7 +23,7 @@ repos:
           --fix
         ]
 -   repo: https://github.com/codespell-project/codespell
-    rev: v2.2.5
+    rev: v2.2.6
     hooks:
     - id: codespell
       additional_dependencies: [ tomli ]

pandas-stubs/core/generic.pyi

@@ -19,6 +19,7 @@ from pandas import Index
 import pandas.core.indexing as indexing
 from pandas.core.series import Series
 import sqlalchemy.engine
+from typing_extensions import Self
 
 from pandas._typing import (
     S1,
@@ -39,7 +40,6 @@ from pandas._typing import (
     IgnoreRaise,
     IndexLabel,
     Level,
-    NDFrameT,
     ReplaceMethod,
     SortKind,
     StorageOptions,
@@ -57,11 +57,11 @@ class NDFrame(indexing.IndexingMixin):
     __hash__: ClassVar[None]  # type: ignore[assignment]
 
     def set_flags(
-        self: NDFrameT,
+        self,
         *,
         copy: bool = ...,
         allows_duplicate_labels: bool | None = ...,
-    ) -> NDFrameT: ...
+    ) -> Self: ...
     @property
     def attrs(self) -> dict[Hashable | None, Any]: ...
     @attrs.setter
@@ -74,19 +74,16 @@ class NDFrame(indexing.IndexingMixin):
     def ndim(self) -> int: ...
     @property
     def size(self) -> int: ...
-    def swapaxes(
-        self, axis1: AxisIndex, axis2: AxisIndex, copy: _bool = ...
-    ) -> NDFrame: ...
-    def droplevel(self, level: Level, axis: AxisIndex = ...) -> NDFrame: ...
+    def droplevel(self, level: Level, axis: AxisIndex = ...) -> Self: ...
     def squeeze(self, axis=...): ...
     def equals(self, other: Series[S1]) -> _bool: ...
-    def __neg__(self: NDFrameT) -> NDFrameT: ...
-    def __pos__(self: NDFrameT) -> NDFrameT: ...
+    def __neg__(self) -> Self: ...
+    def __pos__(self) -> Self: ...
     def __nonzero__(self) -> None: ...
     @final
     def bool(self) -> _bool: ...
-    def __abs__(self) -> NDFrame: ...
-    def __round__(self, decimals: int = ...) -> NDFrame: ...
+    def __abs__(self) -> Self: ...
+    def __round__(self, decimals: int = ...) -> Self: ...
     def keys(self): ...
     def __len__(self) -> int: ...
     def __contains__(self, key) -> _bool: ...
@@ -285,7 +282,7 @@ class NDFrame(indexing.IndexingMixin):
     ) -> _str: ...
     def take(
         self, indices, axis=..., is_copy: _bool | None = ..., **kwargs
-    ) -> NDFrame: ...
+    ) -> Self: ...
     def __delitem__(self, idx: Hashable) -> None: ...
     def get(self, key: object, default: Dtype | None = ...) -> Dtype: ...
     def reindex_like(
@@ -295,7 +292,7 @@ class NDFrame(indexing.IndexingMixin):
         copy: _bool = ...,
         limit=...,
         tolerance=...,
-    ) -> NDFrame: ...
+    ) -> Self: ...
     @overload
     def drop(
         self,
@@ -310,7 +307,7 @@ class NDFrame(indexing.IndexingMixin):
     ) -> None: ...
     @overload
     def drop(
-        self: NDFrame,
+        self,
         labels: Hashable | Sequence[Hashable] = ...,
         *,
         axis: Axis = ...,
@@ -319,10 +316,10 @@ class NDFrame(indexing.IndexingMixin):
         level: Level | None = ...,
         inplace: Literal[False] = ...,
         errors: IgnoreRaise = ...,
-    ) -> NDFrame: ...
+    ) -> Self: ...
     @overload
     def drop(
-        self: NDFrame,
+        self,
         labels: Hashable | Sequence[Hashable] = ...,
         *,
         axis: Axis = ...,
@@ -331,9 +328,9 @@ class NDFrame(indexing.IndexingMixin):
         level: Level | None = ...,
         inplace: _bool = ...,
         errors: IgnoreRaise = ...,
-    ) -> NDFrame | None: ...
-    def add_prefix(self, prefix: _str) -> NDFrame: ...
-    def add_suffix(self, suffix: _str) -> NDFrame: ...
+    ) -> Self | None: ...
+    def add_prefix(self, prefix: _str) -> Self: ...
+    def add_suffix(self, suffix: _str) -> Self: ...
     def sort_index(
         self,
         *,
@@ -352,30 +349,30 @@ class NDFrame(indexing.IndexingMixin):
         like: _str | None = ...,
         regex: _str | None = ...,
         axis=...,
-    ) -> NDFrame: ...
-    def head(self: NDFrameT, n: int = ...) -> NDFrameT: ...
-    def tail(self: NDFrameT, n: int = ...) -> NDFrameT: ...
+    ) -> Self: ...
+    def head(self, n: int = ...) -> Self: ...
+    def tail(self, n: int = ...) -> Self: ...
     def pipe(
         self, func: Callable[..., T] | tuple[Callable[..., T], str], *args, **kwargs
     ) -> T: ...
-    def __finalize__(self, other, method=..., **kwargs) -> NDFrame: ...
+    def __finalize__(self, other, method=..., **kwargs) -> Self: ...
     def __setattr__(self, name: _str, value) -> None: ...
     @property
     def values(self) -> ArrayLike: ...
     @property
     def dtypes(self): ...
-    def copy(self: NDFrameT, deep: _bool = ...) -> NDFrameT: ...
-    def __copy__(self, deep: _bool = ...) -> NDFrame: ...
-    def __deepcopy__(self, memo=...) -> NDFrame: ...
-    def infer_objects(self) -> NDFrame: ...
+    def copy(self, deep: _bool = ...) -> Self: ...
+    def __copy__(self, deep: _bool = ...) -> Self: ...
+    def __deepcopy__(self, memo=...) -> Self: ...
+    def infer_objects(self) -> Self: ...
     def convert_dtypes(
-        self: NDFrameT,
+        self,
         infer_objects: _bool = ...,
         convert_string: _bool = ...,
         convert_integer: _bool = ...,
         convert_boolean: _bool = ...,
         dtype_backend: DtypeBackend = ...,
-    ) -> NDFrameT: ...
+    ) -> Self: ...
     def fillna(
         self,
         value=...,
@@ -403,24 +400,24 @@ class NDFrame(indexing.IndexingMixin):
     def notnull(self) -> NDFrame: ...
     def clip(
         self, lower=..., upper=..., *, axis=..., inplace: _bool = ..., **kwargs
-    ) -> NDFrame: ...
+    ) -> Self: ...
     def asfreq(
         self,
         freq,
         method: FillnaOptions | None = ...,
         how: Literal["start", "end"] | None = ...,
         normalize: _bool = ...,
         fill_value=...,
-    ) -> NDFrame: ...
-    def at_time(self, time, asof: _bool = ..., axis=...) -> NDFrame: ...
+    ) -> Self: ...
+    def at_time(self, time, asof: _bool = ..., axis=...) -> Self: ...
     def between_time(
         self,
         start_time,
         end_time,
         axis=...,
-    ) -> NDFrame: ...
-    def first(self, offset) -> NDFrame: ...
-    def last(self, offset) -> NDFrame: ...
+    ) -> Self: ...
+    def first(self, offset) -> Self: ...
+    def last(self, offset) -> Self: ...
     def rank(
         self,
         axis=...,
@@ -450,13 +447,11 @@ class NDFrame(indexing.IndexingMixin):
         level=...,
         try_cast: _bool = ...,
     ): ...
-    def shift(self, periods=..., freq=..., axis=..., fill_value=...) -> NDFrame: ...
-    def slice_shift(self, periods: int = ..., axis=...) -> NDFrame: ...
-    def tshift(self, periods: int = ..., freq=..., axis=...) -> NDFrame: ...
-    def truncate(
-        self, before=..., after=..., axis=..., copy: _bool = ...
-    ) -> NDFrame: ...
-    def tz_convert(self, tz, axis=..., level=..., copy: _bool = ...) -> NDFrame: ...
+    def shift(self, periods=..., freq=..., axis=..., fill_value=...) -> Self: ...
+    def slice_shift(self, periods: int = ..., axis=...) -> Self: ...
+    def tshift(self, periods: int = ..., freq=..., axis=...) -> Self: ...
+    def truncate(self, before=..., after=..., axis=..., copy: _bool = ...) -> Self: ...
+    def tz_convert(self, tz, axis=..., level=..., copy: _bool = ...) -> Self: ...
     def tz_localize(
         self,
         tz,
@@ -465,8 +460,8 @@ class NDFrame(indexing.IndexingMixin):
         copy: _bool = ...,
         ambiguous=...,
         nonexistent: str = ...,
-    ) -> NDFrame: ...
-    def abs(self) -> NDFrame: ...
+    ) -> Self: ...
+    def abs(self) -> Self: ...
     def describe(
         self,
         percentiles=...,
@@ -476,6 +471,6 @@ class NDFrame(indexing.IndexingMixin):
     ) -> NDFrame: ...
     def pct_change(
         self, periods=..., fill_method=..., limit=..., freq=..., **kwargs
-    ) -> NDFrame: ...
+    ) -> Self: ...
     def first_valid_index(self): ...
     def last_valid_index(self): ...

pandas-stubs/core/series.pyi

@@ -1138,9 +1138,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
     def swapaxes(
         self, axis1: AxisIndex, axis2: AxisIndex, copy: _bool = ...
     ) -> Series[S1]: ...
-    def droplevel(
-        self, level: Level | list[Level], axis: AxisIndex = ...
-    ) -> DataFrame: ...
+    def droplevel(self, level: Level | list[Level], axis: AxisIndex = ...) -> Self: ...
     def pop(self, item: Hashable) -> S1: ...
     def squeeze(self, axis: AxisIndex | None = ...) -> Scalar: ...
     def __abs__(self) -> Series[S1]: ...

tests/test_series.py

@@ -2830,3 +2830,9 @@ def test_series_setitem_na() -> None:
     s2 = df["y"].copy()
     s2.loc[ind] = pd.NaT
     s2.iloc[[0, 2]] = pd.NaT
+
+
+def test_round() -> None:
+    # GH 791
+    check(assert_type(round(pd.DataFrame([])), pd.DataFrame), pd.DataFrame)
+    check(assert_type(round(pd.Series([1], dtype=int)), "pd.Series[int]"), pd.Series)