fix various issues with unknown methods in pandas (#166)

* add concat set_flags Timestamp to deal with unknowns

* fix various unknown declarations

* add partial arg type checking

Author: Dr-Irv

.github/workflows/test.yml

@@ -33,6 +33,7 @@ jobs:
           cp -R partial/ typings
 
       - name: Run tests
-        uses: jakebailey/pyright-action@main
+        uses: gramster/pyright-action@main
         with:
           project: pyrighttestconfig.json
+          warn-partial: true

partial/pandas/_libs/__init__.pyi

@@ -0,0 +1,4 @@
+from .tslibs import NaT as NaT, NaTType as NaTType, OutOfBoundsDatetime as OutOfBoundsDatetime, Period as Period, Timedelta as Timedelta, Timestamp as Timestamp, iNaT as iNaT
+from .interval import Interval as Interval
+
+

partial/pandas/_typing.pyi

@@ -28,6 +28,7 @@ from pandas.core.generic import NDFrame
 from pandas._libs.tslibs import Period, Timedelta, Timestamp
 from pandas.core.arrays import ExtensionArray
 from pandas.core.series import Series
+from pandas.core.frame import DataFrame
 from pandas.core.indexes.base import Index
 
 if sys.version_info >= (3, 8):
@@ -49,6 +50,7 @@ FileOrBuffer = Union[str, Buffer[AnyStr]]
 FilePathOrBuffer = Union["PathLike[str]", FileOrBuffer[AnyStr]]
 
 FrameOrSeries = TypeVar("FrameOrSeries", bound=NDFrame)
+FrameOrSeriesUnion = Union[DataFrame, Series]
 Axis = Union[str, int]
 Label = Optional[Hashable]
 Level = Union[Label, int]

partial/pandas/core/arraylike.pyi

@@ -4,35 +4,35 @@ from pandas._typing import FrameOrSeries, DataFrame
 class OpsMixinProtocol(Protocol): ...
 
 class OpsMixin:
-    def __eq__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
-    def __ne__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
-    def __lt__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
-    def __le__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
-    def __gt__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
-    def __ge__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
+    def __eq__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...
+    def __ne__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...
+    def __lt__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...
+    def __le__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...
+    def __gt__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...
+    def __ge__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...
     # -------------------------------------------------------------
     # Logical Methods
-    def __and__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
-    def __rand__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
-    def __or__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
-    def __ror__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
-    def __xor__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
-    def __rxor__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
+    def __and__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...
+    def __rand__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...
+    def __or__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...
+    def __ror__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...
+    def __xor__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...
+    def __rxor__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...
     # -------------------------------------------------------------
     # Arithmetic Methods
-    def __add__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
-    def __radd__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
-    def __sub__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
-    def __rsub__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
-    def __mul__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
-    def __rmul__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
-    def __truediv__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
-    def __rtruediv__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
-    def __floordiv__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
-    def __rfloordiv__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
-    def __mod__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
-    def __rmod__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
+    def __add__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...
+    def __radd__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...
+    def __sub__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...
+    def __rsub__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...
+    def __mul__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...
+    def __rmul__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...
+    def __truediv__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...
+    def __rtruediv__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...
+    def __floordiv__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...
+    def __rfloordiv__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...
+    def __mod__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...
+    def __rmod__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...
     def __divmod__(self: OpsMixinProtocol, other: DataFrame) -> Tuple[DataFrame, DataFrame]: ...
     def __rdivmod__(self: OpsMixinProtocol, other: DataFrame) -> Tuple[DataFrame, DataFrame]: ...
-    def __pow__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
-    def __rpow__(self: OpsMixinProtocol, other: Any) -> FrameOrSeries: ...
+    def __pow__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...
+    def __rpow__(self: OpsMixinProtocol, other: Any) -> DataFrame: ...

partial/pandas/core/frame.pyi

@@ -233,21 +233,32 @@ class DataFrame(NDFrame, OpsMixin):
         version: int = ...,
         convert_strl: Optional[List[_str]] = ...,
     ) -> None: ...
-    def to_feather(self, path: FileOrBuffer, **kwargs) -> None: ...
+    def to_feather(self, path: FilePathOrBuffer, **kwargs) -> None: ...
     @overload
     def to_markdown(self, buf: Optional[FilePathOrBuffer], mode: Optional[_str] = ..., **kwargs) -> None: ...
     @overload
     def to_markdown(self, mode: Optional[_str] = ..., **kwargs) -> _str: ...
+    @overload
     def to_parquet(
         self,
-        path: _str,
+        path: Optional[FilePathOrBuffer],
         engine: Union[_str, Literal["auto", "pyarrow", "fastparquet"]] = ...,
         compression: Union[_str, Literal["snappy", "gzip", "brotli"]] = ...,
         index: Optional[_bool] = ...,
         partition_cols: Optional[List] = ...,
         **kwargs,
     ) -> None: ...
     @overload
+    def to_parquet(
+        self,
+        path: Optional[FilePathOrBuffer] = ...,
+        engine: Union[_str, Literal["auto", "pyarrow", "fastparquet"]] = ...,
+        compression: Union[_str, Literal["snappy", "gzip", "brotli"]] = ...,
+        index: Optional[_bool] = ...,
+        partition_cols: Optional[List] = ...,
+        **kwargs,
+    ) -> bytes: ...
+    @overload
     def to_html(
         self,
         buf: Optional[FilePathOrBuffer],
@@ -730,6 +741,14 @@ class DataFrame(NDFrame, OpsMixin):
         ignore_index: _bool = ...,
         key: Optional[Callable] = ...,
     ) -> Union[None, DataFrame]: ...
+    def value_counts(
+        self,
+        subset: Optional[Sequence[Hashable]] = ...,
+        normalize: _bool = ...,
+        sort: _bool = ...,
+        ascending: _bool = ...,
+        dropna: _bool = ...,
+    ) -> Series[int]: ...
     def nlargest(
         self,
         n: int,
@@ -744,6 +763,13 @@ class DataFrame(NDFrame, OpsMixin):
     ) -> DataFrame: ...
     def swaplevel(self, i: Level = ..., j: Level = ..., axis: AxisType = ...) -> DataFrame: ...
     def reorder_levels(self, order: List, axis: AxisType = ...) -> DataFrame: ...
+    def compare(
+        self,
+        other: DataFrame,
+        align_axis: Axis = ...,
+        keep_shape: bool = ...,
+        keep_equal: bool = ...,
+    ) -> DataFrame: ...
     def combine(
         self,
         other: DataFrame,

partial/pandas/core/generic.pyi

@@ -12,7 +12,7 @@ from pandas._typing import (
     Scalar as Scalar,
     SeriesAxisType as SeriesAxisType,
     FrameOrSeries as FrameOrSeries,
-    S1 as S1
+    S1 as S1,
 )
 from pandas.core.base import PandasObject as PandasObject, SelectionMixin as SelectionMixin
 from pandas.core.indexes.api import Index as Index
@@ -37,6 +37,7 @@ class NDFrame(PandasObject, SelectionMixin, indexing.IndexingMixin):
         attrs: Optional[Mapping[Optional[Hashable], Any]] = ...,
         fastpath: _bool = ...,
     ) -> None: ...
+    def set_flags(self: FrameOrSeries, *, copy: bool = ..., allows_duplicate_labels: Optional[bool] = ...) -> FrameOrSeries: ...
     @property
     def attrs(self) -> Dict[Optional[Hashable], Any]: ...
     @attrs.setter
@@ -99,7 +100,7 @@ class NDFrame(PandasObject, SelectionMixin, indexing.IndexingMixin):
     def equals(self, other: Series[S1]) -> _bool: ...
     def __neg__(self) -> None: ...
     def __pos__(self) -> None: ...
-    def __invert__(self): ...
+    def __invert__(self) -> NDFrame: ...
     def __nonzero__(self) -> None: ...
     def bool(self) -> _bool: ...
     def __abs__(self) -> NDFrame: ...

partial/pandas/core/groupby/generic.pyi

@@ -24,7 +24,7 @@ def generate_property(name: str, klass: Type[FrameOrSeries]): ...
 def pin_whitelisted_properties(klass: Type[FrameOrSeries], whitelist: FrozenSet[str]): ...
 
 class SeriesGroupBy(GroupBy):
-    def apply(self, func, *args, **kwargs): ...
+    def apply(self, func, *args, **kwargs) -> Series: ...
     def aggregate(self, func=..., *args, **kwargs) -> Series: ...
     agg = aggregate
     def transform(self, func, *args, **kwargs): ...
@@ -71,6 +71,7 @@ class SeriesGroupBy(GroupBy):
     def nth(self, n: Union[int, Sequence[int]], dropna: Optional[str] = ...) -> Series[S1]: ...
 
 class DataFrameGroupBy(GroupBy):
+    def apply(self, func, *args, **kwargs) -> DataFrame: ...
     @overload
     def aggregate(self, arg: str, *args, **kwargs) -> DataFrame: ...
     @overload
@@ -85,6 +86,7 @@ class DataFrameGroupBy(GroupBy):
     def agg(self, arg: F, *args, **kwargs) -> DataFrame: ...
     def transform(self, func, *args, **kwargs): ...
     def filter(self, func: Callable, dropna: bool = ..., *args, **kwargs) -> DataFrame: ...
+    def nunique(self, dropna: bool = ...) -> DataFrame: ...
     @overload
     def __getitem__(self, item: str) -> SeriesGroupBy: ...
     @overload

partial/pandas/core/groupby/groupby.pyi

@@ -1,4 +1,10 @@
-from pandas._typing import FrameOrSeries as FrameOrSeries, Scalar as Scalar, AxisType as AxisType, KeysArgType
+from pandas._typing import (
+    FrameOrSeries as FrameOrSeries,
+    FrameOrSeriesUnion as FrameOrSeriesUnion,
+    Scalar as Scalar,
+    AxisType as AxisType,
+    KeysArgType,
+)
 from pandas.core.base import PandasObject as PandasObject, SelectionMixin as SelectionMixin
 from pandas.core.frame import DataFrame as DataFrame
 from pandas.core.generic import NDFrame as NDFrame
@@ -7,11 +13,10 @@ from pandas.core.indexes.api import Index as Index
 from pandas.core.series import Series as Series
 from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union
 
-
 class GroupByPlot(PandasObject):
     def __init__(self, groupby) -> None: ...
     def __call__(self, *args, **kwargs): ...
-    def __getattr__(self, name: str) : ...
+    def __getattr__(self, name: str): ...
 
 class _GroupBy(PandasObject, SelectionMixin):
     level = ...
@@ -26,66 +31,103 @@ class _GroupBy(PandasObject, SelectionMixin):
     axis = ...
     grouper = ...
     exclusions = ...
-    def __init__(self, obj: NDFrame, keys: Optional[KeysArgType]=..., axis: int=..., level=..., grouper: Optional[ops.BaseGrouper]=..., exclusions=..., selection=..., as_index: bool=..., sort: bool=..., group_keys: bool=..., squeeze: bool=..., observed: bool=..., mutated: bool=...) -> None: ...
+    def __init__(
+        self,
+        obj: NDFrame,
+        keys: Optional[KeysArgType] = ...,
+        axis: int = ...,
+        level=...,
+        grouper: Optional[ops.BaseGrouper] = ...,
+        exclusions=...,
+        selection=...,
+        as_index: bool = ...,
+        sort: bool = ...,
+        group_keys: bool = ...,
+        squeeze: bool = ...,
+        observed: bool = ...,
+        mutated: bool = ...,
+    ) -> None: ...
     def __len__(self) -> int: ...
     @property
     def groups(self) -> Dict[str, str]: ...
     @property
     def ngroups(self): ...
     @property
     def indices(self) -> Dict[str, Index]: ...
-    def __getattr__(self, attr: str) : ...
-    def pipe(self, func: Callable, *args, **kwargs) : ...
+    def __getattr__(self, attr: str): ...
+    def pipe(self, func: Callable, *args, **kwargs): ...
     plot = ...
     def get_group(self, name, obj: Optional[DataFrame] = ...) -> DataFrame: ...
     def __iter__(self) -> Generator[Tuple[str, Any], None, None]: ...
-    def apply(self, func: Callable, *args, **kwargs) -> FrameOrSeries: ...
+    def apply(self, func: Callable, *args, **kwargs) -> FrameOrSeriesUnion: ...
     def transform(self, func: Callable, *args, **kwargs): ...
 
-
 class GroupBy(_GroupBy):
-    def any(self, skipna: bool=...) -> bool: ...
-    def all(self, skipna: bool=...) -> bool: ...
-    def count(self) -> FrameOrSeries: ...
-    def mean(self, **kwargs) -> FrameOrSeries: ...
-    def median(self, **kwargs) -> FrameOrSeries: ...
-    def std(self, ddof: int = ...) -> FrameOrSeries: ...
-    def var(self, ddof: int = ...) -> FrameOrSeries: ...
-    def sem(self, ddof: int = ...) -> FrameOrSeries: ...
+    def any(self, skipna: bool = ...) -> bool: ...
+    def all(self, skipna: bool = ...) -> bool: ...
+    def count(self) -> FrameOrSeriesUnion: ...
+    def mean(self, **kwargs) -> FrameOrSeriesUnion: ...
+    def median(self, **kwargs) -> FrameOrSeriesUnion: ...
+    def std(self, ddof: int = ...) -> FrameOrSeriesUnion: ...
+    def var(self, ddof: int = ...) -> FrameOrSeriesUnion: ...
+    def sem(self, ddof: int = ...) -> FrameOrSeriesUnion: ...
     def size(self) -> Series: ...
     def ohlc(self) -> DataFrame: ...
-    def describe(self, **kwargs) -> FrameOrSeries: ...
+    def describe(self, **kwargs) -> FrameOrSeriesUnion: ...
     def resample(self, rule, *args, **kwargs): ...
     def rolling(self, *args, **kwargs): ...
     def expanding(self, *args, **kwargs): ...
     def pad(self, limit: Optional[int] = ...): ...
-    def ffill(self, limit: Optional[int] = ...) -> FrameOrSeries: ...
-    def backfill(self, limit: Optional[int] = ...) -> FrameOrSeries: ...
-    def bfill(self, limit: Optional[int] = ...) -> FrameOrSeries: ...
-    def nth(self, n: Union[int, List[int]], dropna: Optional[str] = ...) -> FrameOrSeries: ...
-    def quantile(self, q=..., interpolation: str=...) : ...
+    def ffill(self, limit: Optional[int] = ...) -> FrameOrSeriesUnion: ...
+    def backfill(self, limit: Optional[int] = ...) -> FrameOrSeriesUnion: ...
+    def bfill(self, limit: Optional[int] = ...) -> FrameOrSeriesUnion: ...
+    def nth(self, n: Union[int, List[int]], dropna: Optional[str] = ...) -> FrameOrSeriesUnion: ...
+    def quantile(self, q=..., interpolation: str = ...): ...
     def ngroup(self, ascending: bool = ...) -> Series: ...
     def cumcount(self, ascending: bool = ...) -> Series: ...
     def rank(
-        self, method: str = ..., ascending: bool = ..., na_option: str = ..., pct: bool = ..., axis: int = ...,
+        self,
+        method: str = ...,
+        ascending: bool = ...,
+        na_option: str = ...,
+        pct: bool = ...,
+        axis: int = ...,
     ) -> DataFrame: ...
-    def cummax(self, axis: AxisType = ..., **kwargs) -> FrameOrSeries: ...
-    def cummin(self, axis: AxisType = ..., **kwargs) -> FrameOrSeries: ...
-    def cumprod(self, axis: AxisType = ..., **kwargs) -> FrameOrSeries: ...
-    def cumsum(self, axis: AxisType = ..., **kwargs) -> FrameOrSeries: ...
-    def shift(self, periods: int = ..., freq = ..., axis: int = ..., fill_value = ...): ...
+    def cummax(self, axis: AxisType = ..., **kwargs) -> FrameOrSeriesUnion: ...
+    def cummin(self, axis: AxisType = ..., **kwargs) -> FrameOrSeriesUnion: ...
+    def cumprod(self, axis: AxisType = ..., **kwargs) -> FrameOrSeriesUnion: ...
+    def cumsum(self, axis: AxisType = ..., **kwargs) -> FrameOrSeriesUnion: ...
+    def shift(self, periods: int = ..., freq=..., axis: int = ..., fill_value=...): ...
     def pct_change(
-        self, periods: int = ..., fill_method: str = ..., limit = ..., freq = ..., axis: AxisType = ...,
-    ) -> FrameOrSeries: ...
-    def head(self, n: int = ...) -> FrameOrSeries: ...
-    def tail(self, n: int = ...) -> FrameOrSeries: ...
+        self,
+        periods: int = ...,
+        fill_method: str = ...,
+        limit=...,
+        freq=...,
+        axis: AxisType = ...,
+    ) -> FrameOrSeriesUnion: ...
+    def head(self, n: int = ...) -> FrameOrSeriesUnion: ...
+    def tail(self, n: int = ...) -> FrameOrSeriesUnion: ...
     # Surplus methodss from original pylance stubs; should they go away?
-    def first(self, **kwargs) -> FrameOrSeries: ...
-    def last(self, **kwargs) -> FrameOrSeries: ...
-    def max(self, **kwargs) -> FrameOrSeries: ...
-    def min(self, **kwargs) -> FrameOrSeries: ...
-    def prod(self, **kwargs) -> FrameOrSeries: ...
-    def sum(self, **kwargs) -> FrameOrSeries: ...
-
+    def first(self, **kwargs) -> FrameOrSeriesUnion: ...
+    def last(self, **kwargs) -> FrameOrSeriesUnion: ...
+    def max(self, **kwargs) -> FrameOrSeriesUnion: ...
+    def min(self, **kwargs) -> FrameOrSeriesUnion: ...
+    def prod(self, **kwargs) -> FrameOrSeriesUnion: ...
+    def sum(self, **kwargs) -> FrameOrSeriesUnion: ...
 
-def get_groupby(obj: NDFrame, by: Optional[KeysArgType]=..., axis: int=..., level=..., grouper: Optional[ops.BaseGrouper]=..., exclusions=..., selection=..., as_index: bool=..., sort: bool=..., group_keys: bool=..., squeeze: bool=..., observed: bool=..., mutated: bool=...) -> GroupBy: ...
+def get_groupby(
+    obj: NDFrame,
+    by: Optional[KeysArgType] = ...,
+    axis: int = ...,
+    level=...,
+    grouper: Optional[ops.BaseGrouper] = ...,
+    exclusions=...,
+    selection=...,
+    as_index: bool = ...,
+    sort: bool = ...,
+    group_keys: bool = ...,
+    squeeze: bool = ...,
+    observed: bool = ...,
+    mutated: bool = ...,
+) -> GroupBy: ...