__neg__ methods and bugfix (#253)

* Added __neg__ to Series, DataFrame and Index. Fixed bug in Index.difference

* Moved to NDFrameT and IndexT

* sorted imports

Co-authored-by: Amotz <amotz@nonaya.bis>

Author: amotzop

pandas-stubs/_typing.pyi

@@ -169,6 +169,8 @@ IndexingInt = Union[
 # passed in, a DataFrame is always returned.
 NDFrameT = TypeVar("NDFrameT", bound=NDFrame)
 
+IndexT = TypeVar("IndexT", bound=Index)
+
 # Interval closed type
 
 IntervalClosedType = Literal["left", "right", "both", "neither"]

pandas-stubs/core/generic.pyi

@@ -29,6 +29,7 @@ from pandas._typing import (
     HDFCompLib,
     IgnoreRaise,
     Level,
+    NDFrameT,
     ReplaceMethod,
     Scalar,
     SeriesAxisType,
@@ -69,8 +70,8 @@ class NDFrame(PandasObject, indexing.IndexingMixin):
     def pop(self, item: _str) -> NDFrame: ...
     def squeeze(self, axis=...): ...
     def equals(self, other: Series[S1]) -> _bool: ...
-    def __neg__(self) -> None: ...
-    def __pos__(self) -> None: ...
+    def __neg__(self: NDFrameT) -> NDFrameT: ...
+    def __pos__(self: NDFrameT) -> NDFrameT: ...
     def __nonzero__(self) -> None: ...
     def bool(self) -> _bool: ...
     def __abs__(self) -> NDFrame: ...

pandas-stubs/core/indexes/base.pyi

@@ -27,6 +27,7 @@ from pandas._typing import (
     Dtype,
     DtypeArg,
     DtypeObj,
+    IndexT,
     Label,
     Level,
     NaPosition,
@@ -142,11 +143,12 @@ class Index(IndexOpsMixin, PandasObject):
     def __and__(self, other) -> Index: ...
     def __or__(self, other) -> Index: ...
     def __xor__(self, other) -> Index: ...
+    def __neg__(self: IndexT) -> IndexT: ...
     def __nonzero__(self) -> None: ...
     __bool__ = ...
     def union(self, other: list[T1] | Index, sort=...) -> Index: ...
     def intersection(self, other: list[T1] | Index, sort: bool = ...) -> Index: ...
-    def difference(self, other: list[T1 | None] | Index) -> Index: ...
+    def difference(self, other: list | Index) -> Index: ...
     def symmetric_difference(
         self, other: list[T1] | Index, result_name=..., sort=...
     ) -> Index: ...

tests/test_frame.py

@@ -1685,3 +1685,15 @@ def func() -> MyDataFrame[int]:
         return MyDataFrame[int]({"foo": [1, 2, 3]})
 
     func()
+
+
+def test_neg() -> None:
+    # GH 253
+    df = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]})
+    check(assert_type(-df, pd.DataFrame), pd.DataFrame)
+
+
+def test_pos() -> None:
+    # GH 253
+    df = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]})
+    check(assert_type(+df, pd.DataFrame), pd.DataFrame)

tests/test_indexes.py

@@ -9,7 +9,6 @@
 
 
 def test_index_unique() -> None:
-
     df = pd.DataFrame({"x": [1, 2, 3, 4]}, index=pd.Index([1, 2, 3, 2]))
     ind = df.index
     check(assert_type(ind, pd.Index), pd.Index)
@@ -67,7 +66,9 @@ def test_column_contains() -> None:
 def test_difference_none() -> None:
     # https://github.com/pandas-dev/pandas-stubs/issues/17
     ind = pd.Index([1, 2, 3])
-    check(assert_type(ind.difference([1, None]), "pd.Index"), pd.Index, int)
+    check(assert_type(ind.difference([1, None]), pd.Index), pd.Index)
+    # GH 253
+    check(assert_type(ind.difference([1]), pd.Index), pd.Index)
 
 
 def test_str_split() -> None:
@@ -87,3 +88,9 @@ def test_index_dropna():
 
     check(assert_type(midx.dropna(how="all"), pd.MultiIndex), pd.MultiIndex)
     check(assert_type(midx.dropna(how="any"), pd.MultiIndex), pd.MultiIndex)
+
+
+def test_index_neg():
+    # GH 253
+    idx = pd.Index([1, 2])
+    check(assert_type(-idx, pd.Index), pd.Index)

tests/test_series.py

@@ -1124,3 +1124,11 @@ def test_resample() -> None:
     check(assert_type(df.resample("2T").sem(), pd.Series), pd.Series)
     check(assert_type(df.resample("2T").median(), pd.Series), pd.Series)
     check(assert_type(df.resample("2T").ohlc(), pd.DataFrame), pd.DataFrame)
+
+
+def test_neg() -> None:
+    # GH 253
+    sr = pd.Series([1, 2, 3])
+    sr_int = pd.Series([1, 2, 3], dtype=int)
+    check(assert_type(-sr, pd.Series), pd.Series)
+    check(assert_type(-sr_int, "pd.Series[int]"), pd.Series, int)