diff --git a/pandas-stubs/_libs/missing.pyi b/pandas-stubs/_libs/missing.pyi index 2c948529..c06bf361 100644 --- a/pandas-stubs/_libs/missing.pyi +++ b/pandas-stubs/_libs/missing.pyi @@ -1,46 +1,243 @@ +from typing import ( + Any, + Callable, + Literal, + overload, +) + +from pandas import ( + Index, + Series, +) +from pandas.core.arrays.boolean import BooleanArray from typing_extensions import Self class NAType: - def __new__(cls, *args, **kwargs) -> Self: ... + def __new__(cls, *args: Any, **kwargs: Any) -> Self: ... def __format__(self, format_spec: str) -> str: ... - def __bool__(self) -> None: ... def __hash__(self) -> int: ... def __reduce__(self) -> str: ... - def __add__(self, other) -> NAType: ... - def __radd__(self, other) -> NAType: ... - def __sub__(self, other) -> NAType: ... - def __rsub__(self, other) -> NAType: ... - def __mul__(self, other) -> NAType: ... - def __rmul__(self, other) -> NAType: ... - def __matmul__(self, other) -> NAType: ... - def __rmatmul__(self, other) -> NAType: ... - def __truediv__(self, other) -> NAType: ... - def __rtruediv__(self, other) -> NAType: ... - def __floordiv__(self, other) -> NAType: ... - def __rfloordiv__(self, other) -> NAType: ... - def __mod__(self, other) -> NAType: ... - def __rmod__(self, other) -> NAType: ... - def __divmod__(self, other) -> NAType: ... - def __rdivmod__(self, other) -> NAType: ... - def __eq__(self, other) -> bool: ... - def __ne__(self, other) -> bool: ... - def __le__(self, other) -> bool: ... - def __lt__(self, other) -> bool: ... - def __gt__(self, other) -> bool: ... - def __ge__(self, other) -> bool: ... - def __neg__(self, other) -> NAType: ... - def __pos__(self, other) -> NAType: ... - def __abs__(self, other) -> NAType: ... - def __invert__(self, other) -> NAType: ... - def __pow__(self, other) -> NAType: ... - def __rpow__(self, other) -> NAType: ... - def __and__(self, other) -> NAType | None: ... - __rand__ = __and__ - def __or__(self, other) -> bool | NAType: ... - __ror__ = __or__ - def __xor__(self, other) -> NAType: ... - __rxor__ = __xor__ + @overload + def __add__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self, other: Series, / + ) -> Series: ... + @overload + def __add__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] + @overload + def __add__(self, other: object, /) -> NAType: ... + @overload + def __radd__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self, other: Series, / + ) -> Series: ... + @overload + def __radd__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] + @overload + def __radd__(self, other: object, /) -> NAType: ... + @overload + def __sub__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self, other: Series, / + ) -> Series: ... + @overload + def __sub__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] + @overload + def __sub__(self, other: object, /) -> NAType: ... + @overload + def __rsub__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self, other: Series, / + ) -> Series: ... + @overload + def __rsub__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] + @overload + def __rsub__(self, other: object, /) -> NAType: ... + @overload + def __mul__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self, other: Series, / + ) -> Series: ... + @overload + def __mul__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] + @overload + def __mul__(self, other: object, /) -> NAType: ... + @overload + def __rmul__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self, other: Series, / + ) -> Series: ... + @overload + def __rmul__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] + @overload + def __rmul__(self, other: object, /) -> NAType: ... + def __matmul__(self, other: object, /) -> NAType: ... + def __rmatmul__(self, other: object, /) -> NAType: ... + @overload + def __truediv__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self, other: Series, / + ) -> Series: ... + @overload + def __truediv__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] + @overload + def __truediv__(self, other: object, /) -> NAType: ... + @overload + def __rtruediv__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self, other: Series, / + ) -> Series: ... + @overload + def __rtruediv__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] + @overload + def __rtruediv__(self, other: object, /) -> NAType: ... + @overload + def __floordiv__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self, other: Series, / + ) -> Series: ... + @overload + def __floordiv__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] + @overload + def __floordiv__(self, other: object, /) -> NAType: ... + @overload + def __rfloordiv__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self, other: Series, / + ) -> Series: ... + @overload + def __rfloordiv__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] + @overload + def __rfloordiv__(self, other: object, /) -> NAType: ... + @overload + def __mod__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self, other: Series, / + ) -> Series: ... + @overload + def __mod__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] + @overload + def __mod__(self, other: object, /) -> NAType: ... + @overload + def __rmod__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self, other: Series, / + ) -> Series: ... + @overload + def __rmod__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] + @overload + def __rmod__(self, other: object, /) -> NAType: ... + @overload + def __divmod__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self, other: Series, / + ) -> tuple[Series, Series]: ... + @overload + def __divmod__(self, other: Index, /) -> tuple[Index, Index]: ... # type: ignore[overload-overlap] + @overload + def __divmod__(self, other: object, /) -> tuple[NAType, NAType]: ... + @overload + def __rdivmod__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self, other: Series, / + ) -> tuple[Series, Series]: ... + @overload + def __rdivmod__(self, other: Index, /) -> tuple[Index, Index]: ... # type: ignore[overload-overlap] + @overload + def __rdivmod__(self, other: object, /) -> tuple[NAType, NAType]: ... + @overload # type: ignore[override] + def __eq__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self, other: Series, / + ) -> Series[bool]: ... + @overload + def __eq__(self, other: Index, /) -> BooleanArray: ... # type: ignore[overload-overlap] + @overload + def __eq__( # pyright: ignore[reportIncompatibleMethodOverride] + self, other: object, / + ) -> NAType: ... + @overload # type: ignore[override] + def __ne__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self, other: Series, / + ) -> Series[bool]: ... + @overload + def __ne__(self, other: Index, /) -> BooleanArray: ... # type: ignore[overload-overlap] + @overload + def __ne__( # pyright: ignore[reportIncompatibleMethodOverride] + self, other: object, / + ) -> NAType: ... + @overload + def __le__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self, other: Series, / + ) -> Series[bool]: ... + @overload + def __le__(self, other: Index, /) -> BooleanArray: ... # type: ignore[overload-overlap] + @overload + def __le__(self, other: object, /) -> NAType: ... + @overload + def __lt__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self, other: Series, / + ) -> Series[bool]: ... + @overload + def __lt__(self, other: Index, /) -> BooleanArray: ... # type: ignore[overload-overlap] + @overload + def __lt__(self, other: object, /) -> NAType: ... + @overload + def __gt__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self, other: Series, / + ) -> Series[bool]: ... + @overload + def __gt__(self, other: Index, /) -> BooleanArray: ... # type: ignore[overload-overlap] + @overload + def __gt__(self, other: object, /) -> NAType: ... + @overload + def __ge__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self, other: Series, / + ) -> Series[bool]: ... + @overload + def __ge__(self, other: Index, /) -> BooleanArray: ... # type: ignore[overload-overlap] + @overload + def __ge__(self, other: object, /) -> NAType: ... + def __neg__(self) -> NAType: ... + def __pos__(self) -> NAType: ... + def __abs__(self) -> NAType: ... + def __invert__(self) -> NAType: ... + @overload + def __pow__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self, other: Series, / + ) -> Series: ... + @overload + def __pow__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] + @overload + def __pow__(self, other: object, /) -> NAType: ... + @overload + def __rpow__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self, other: Series, / + ) -> Series: ... + @overload + def __rpow__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] + @overload + def __rpow__(self, other: object, /) -> NAType: ... + @overload + def __and__(self, other: Literal[False], /) -> Literal[False]: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + @overload + def __and__(self, other: bool | NAType, /) -> NAType: ... + @overload + def __rand__(self, other: Literal[False], /) -> Literal[False]: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + @overload + def __rand__(self, other: bool, /) -> NAType: ... + @overload + def __or__(self, other: Literal[True], /) -> Literal[True]: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + @overload + def __or__(self, other: bool | NAType, /) -> NAType: ... + @overload + def __ror__(self, other: Literal[True], /) -> Literal[True]: ... # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + @overload + def __ror__(self, other: bool | NAType, /) -> NAType: ... + @overload + def __xor__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self, other: Series, / + ) -> Series: ... + @overload + def __xor__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] + @overload + def __xor__(self, other: object, /) -> NAType: ... + @overload + def __rxor__( # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload] + self, other: Series, / + ) -> Series: ... + @overload + def __rxor__(self, other: Index, /) -> Index: ... # type: ignore[overload-overlap] + @overload + def __rxor__(self, other: object, /) -> NAType: ... __array_priority__: int - def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): ... + def __array_ufunc__( + self, ufunc: Callable[..., Any], method: str, *inputs: Any, **kwargs: Any + ) -> Any: ... NA: NAType = ... diff --git a/tests/test_natype.py b/tests/test_natype.py new file mode 100644 index 00000000..1d80c185 --- /dev/null +++ b/tests/test_natype.py @@ -0,0 +1,186 @@ +from typing import ( + Literal, +) + +import pandas as pd +from pandas.api.typing import NAType +from pandas.core.arrays.boolean import BooleanArray +from typing_extensions import assert_type + +from tests import check + + +def test_arithmetic() -> None: + na = pd.NA + + s_int = pd.Series([1, 2, 3], dtype="Int64") + idx_int: pd.Index[int] = pd.Index([1, 2, 3], dtype="Int64") + + # __add__ + check(assert_type(na + s_int, pd.Series), pd.Series) + check(assert_type(na + idx_int, pd.Index), pd.Index) + check(assert_type(na + 1, NAType), NAType) + + # __radd__ + check(assert_type(s_int + na, pd.Series), pd.Series) + # requires orthogonal fix: https://github.com/pandas-dev/pandas-stubs/issues/1347 + # check(assert_type(idx_int + na, pd.Index), pd.Index) + check(assert_type(1 + na, NAType), NAType) + + # __sub__ + check(assert_type(na - s_int, pd.Series), pd.Series) + check(assert_type(na - idx_int, pd.Index), pd.Index) + check(assert_type(na - 1, NAType), NAType) + + # __rsub__ + check(assert_type(s_int - na, pd.Series), pd.Series) + # requires orthogonal fix: https://github.com/pandas-dev/pandas-stubs/issues/1347 + # check(assert_type(idx_int - na, pd.Index), pd.Index) + check(assert_type(1 - na, NAType), NAType) + + # __mul__ + check(assert_type(na * s_int, pd.Series), pd.Series) + check(assert_type(na * idx_int, pd.Index), pd.Index) + check(assert_type(na * 1, NAType), NAType) + + # __rmul__ + check(assert_type(s_int * na, pd.Series), pd.Series) + # requires orthogonal fix: https://github.com/pandas-dev/pandas-stubs/issues/1347 + # check(assert_type(idx_int * na, pd.Index), pd.Index) + check(assert_type(1 * na, NAType), NAType) + + # __matmul__ + check(assert_type(na @ 1, NAType), NAType) + + # __rmatmul__ + check(assert_type(1 @ na, NAType), NAType) + + # __truediv__ + check(assert_type(na / s_int, pd.Series), pd.Series) + check(assert_type(na / idx_int, pd.Index), pd.Index) + check(assert_type(na / 1, NAType), NAType) + + # __rtruediv__ + check(assert_type(s_int / na, pd.Series), pd.Series) + check(assert_type(idx_int / na, pd.Index), pd.Index) + check(assert_type(1 / na, NAType), NAType) + + # __floordiv__ + check(assert_type(na // s_int, pd.Series), pd.Series) + check(assert_type(na // idx_int, pd.Index), pd.Index) + check(assert_type(na // 1, NAType), NAType) + + # __rfloordiv__ + check(assert_type(s_int // na, pd.Series), pd.Series) + check(assert_type(idx_int // na, pd.Index), pd.Index) + check(assert_type(1 // na, NAType), NAType) + + # __mod__ + check(assert_type(na % s_int, pd.Series), pd.Series) + check(assert_type(na % idx_int, pd.Index), pd.Index) + check(assert_type(na % 1, NAType), NAType) + + # __rmod__ + check(assert_type(s_int % na, pd.Series), pd.Series) + # requires orthogonal fix: https://github.com/pandas-dev/pandas-stubs/issues/1347 + # check(assert_type(idx_int % na, pd.Index), pd.Index) + check(assert_type(1 % na, NAType), NAType) + + # __divmod__ + # bug upstream: https://github.com/pandas-dev/pandas/issues/62196 + # check( + # assert_type( + # divmod(na, s_int), + # tuple[pd.Series, pd.Series], + # ), + # tuple, + # ) + # check( + # assert_type( + # # pyright bug: https://github.com/microsoft/pyright/issues/10849, will + # # require `pyright: ignore[reportAssertTypeFailure]` until fixed. + # divmod(na, idx_int), + # tuple[pd.Index, pd.Index], + # ), + # tuple, + # ) + check(assert_type(divmod(na, 1), tuple[NAType, NAType]), tuple) + + # __rdivmod__ + # bug upstream: https://github.com/pandas-dev/pandas/issues/62196 + # check( + # assert_type(divmod(s_int, na), tuple[pd.Series, pd.Series]), + # tuple, + # ) + # https://github.com/pandas-dev/pandas-stubs/issues/1347 + # check( + # assert_type(divmod(idx_int, na), tuple[pd.Index, pd.Index]), + # tuple, + # ) + check(assert_type(divmod(1, na), tuple[NAType, NAType]), tuple) + + # __eq__ + check(assert_type(na == s_int, "pd.Series[bool]"), pd.Series) + check(assert_type(na == idx_int, BooleanArray), BooleanArray) + check(assert_type(na == 1, NAType), NAType) + + # __ne__ + check(assert_type(na != s_int, "pd.Series[bool]"), pd.Series) + check(assert_type(na != idx_int, BooleanArray), BooleanArray) + check(assert_type(na != 1, NAType), NAType) + + # __le__ + check(assert_type(na <= s_int, "pd.Series[bool]"), pd.Series) + check(assert_type(na <= idx_int, BooleanArray), BooleanArray) + check(assert_type(na <= 1, NAType), NAType) + + # __lt__ + check(assert_type(na < s_int, "pd.Series[bool]"), pd.Series) + check(assert_type(na < idx_int, BooleanArray), BooleanArray) + check(assert_type(na < 1, NAType), NAType) + + # __gt__ + check(assert_type(na > s_int, "pd.Series[bool]"), pd.Series) + check(assert_type(na > idx_int, BooleanArray), BooleanArray) + check(assert_type(na > 1, NAType), NAType) + + # __ge__ + check(assert_type(na >= s_int, "pd.Series[bool]"), pd.Series) + check(assert_type(na >= idx_int, BooleanArray), BooleanArray) + check(assert_type(na >= 1, NAType), NAType) + + # __pow__ + check(assert_type(na**s_int, pd.Series), pd.Series) + check(assert_type(na**idx_int, pd.Index), pd.Index) + check(assert_type(na**2, NAType), NAType) + + # __rpow__ + check(assert_type(s_int**na, pd.Series), pd.Series) + # requires orthogonal fix: https://github.com/pandas-dev/pandas-stubs/issues/1347 + # check(assert_type(idx_int**na, pd.Index), pd.Index) + check(assert_type(2**na, NAType), NAType) + + # __and__ + check(assert_type(na & False, Literal[False]), bool) + check(assert_type(na & True, NAType), NAType) + check(assert_type(na & na, NAType), NAType) + + # __rand__ + check(assert_type(False & na, Literal[False]), bool) + check(assert_type(True & na, NAType), NAType) + + # __or__ + check(assert_type(na | False, NAType), NAType) + check(assert_type(na | True, Literal[True]), bool) + + # __ror__ + check(assert_type(False | na, NAType), NAType) + check(assert_type(True | na, Literal[True]), bool) + + # __xor__ + check(assert_type(na ^ s_int, pd.Series), pd.Series) + check(assert_type(na ^ idx_int, pd.Index), pd.Index) + + # rxor + check(assert_type(s_int ^ na, pd.Series), pd.Series) + check(assert_type(idx_int ^ na, pd.Index), pd.Index)