TYP: fix reportInvalidTypeVarUse from pyright (#42367)

Author: twoertwein

pandas/_libs/tslibs/dtypes.pyi

@@ -9,7 +9,7 @@ class PeriodDtypeBase:
     _dtype_code: int  # PeriodDtypeCode
 
     # actually __cinit__
-    def __new__(self, code: int): ...
+    def __new__(cls, code: int): ...
     def freq_group_code(self) -> int: ...
     def date_offset(self) -> BaseOffset: ...
     @classmethod

pandas/_testing/_io.py

@@ -4,15 +4,13 @@
 from functools import wraps
 import gzip
 from typing import (
+    TYPE_CHECKING,
     Any,
     Callable,
 )
 import zipfile
 
-from pandas._typing import (
-    FilePathOrBuffer,
-    FrameOrSeries,
-)
+from pandas._typing import FilePathOrBuffer
 from pandas.compat import (
     get_lzma_file,
     import_lzma,
@@ -24,6 +22,12 @@
 
 from pandas.io.common import urlopen
 
+if TYPE_CHECKING:
+    from pandas import (
+        DataFrame,
+        Series,
+    )
+
 _RAISE_NETWORK_ERROR_DEFAULT = False
 
 lzma = import_lzma()
@@ -272,7 +276,9 @@ def can_connect(url, error_classes=None):
 # File-IO
 
 
-def round_trip_pickle(obj: Any, path: FilePathOrBuffer | None = None) -> FrameOrSeries:
+def round_trip_pickle(
+    obj: Any, path: FilePathOrBuffer | None = None
+) -> DataFrame | Series:
     """
     Pickle an object and then read it again.
 

pandas/core/apply.py

@@ -1338,7 +1338,7 @@ def _make_unique_kwarg_list(
 
 
 def relabel_result(
-    result: FrameOrSeries,
+    result: DataFrame | Series,
     func: dict[str, list[Callable | str]],
     columns: Iterable[Hashable],
     order: Iterable[int],

pandas/core/arrays/_mixins.py

@@ -323,7 +323,7 @@ def __repr__(self) -> str:
     # ------------------------------------------------------------------------
     # __array_function__ methods
 
-    def putmask(self: NDArrayBackedExtensionArrayT, mask: np.ndarray, value) -> None:
+    def putmask(self, mask: np.ndarray, value) -> None:
         """
         Analogue to np.putmask(self, mask, value)
 

pandas/core/arrays/interval.py

@@ -934,9 +934,7 @@ def copy(self: IntervalArrayT) -> IntervalArrayT:
     def isna(self) -> np.ndarray:
         return isna(self._left)
 
-    def shift(
-        self: IntervalArrayT, periods: int = 1, fill_value: object = None
-    ) -> IntervalArray:
+    def shift(self, periods: int = 1, fill_value: object = None) -> IntervalArray:
         if not len(self) or periods == 0:
             return self.copy()
 

pandas/core/generic.py

@@ -709,13 +709,11 @@ def set_axis(
         ...
 
     @overload
-    def set_axis(
-        self: FrameOrSeries, labels, axis: Axis, inplace: Literal[True]
-    ) -> None:
+    def set_axis(self, labels, axis: Axis, inplace: Literal[True]) -> None:
         ...
 
     @overload
-    def set_axis(self: FrameOrSeries, labels, *, inplace: Literal[True]) -> None:
+    def set_axis(self, labels, *, inplace: Literal[True]) -> None:
         ...
 
     @overload
@@ -2703,10 +2701,12 @@ def to_hdf(
         """
         from pandas.io import pytables
 
+        # Argument 3 to "to_hdf" has incompatible type "NDFrame"; expected
+        # "Union[DataFrame, Series]" [arg-type]
         pytables.to_hdf(
             path_or_buf,
             key,
-            self,
+            self,  # type: ignore[arg-type]
             mode=mode,
             complevel=complevel,
             complib=complib,
@@ -3850,7 +3850,7 @@ def _slice(self: FrameOrSeries, slobj: slice, axis=0) -> FrameOrSeries:
         return result
 
     @final
-    def _set_is_copy(self, ref: FrameOrSeries, copy: bool_t = True) -> None:
+    def _set_is_copy(self, ref: NDFrame, copy: bool_t = True) -> None:
         if not copy:
             self._is_copy = None
         else:
@@ -11019,13 +11019,11 @@ def ewm(
         adjust: bool_t = True,
         ignore_na: bool_t = False,
         axis: Axis = 0,
-        times: str | np.ndarray | FrameOrSeries | None = None,
+        times: str | np.ndarray | DataFrame | Series | None = None,
         method: str = "single",
     ) -> ExponentialMovingWindow:
         axis = self._get_axis_number(axis)
-        # error: Value of type variable "FrameOrSeries" of "ExponentialMovingWindow"
-        # cannot be "object"
-        return ExponentialMovingWindow(  # type: ignore[type-var]
+        return ExponentialMovingWindow(
             self,
             com=com,
             span=span,

pandas/core/groupby/generic.py

@@ -1523,7 +1523,7 @@ def _wrap_agged_manager(self, mgr: Manager2D) -> DataFrame:
         #  values converted
         return self._reindex_output(result)._convert(datetime=True)
 
-    def _iterate_column_groupbys(self, obj: FrameOrSeries):
+    def _iterate_column_groupbys(self, obj: DataFrame | Series):
         for i, colname in enumerate(obj.columns):
             yield colname, SeriesGroupBy(
                 obj.iloc[:, i],
@@ -1532,7 +1532,7 @@ def _iterate_column_groupbys(self, obj: FrameOrSeries):
                 exclusions=self.exclusions,
             )
 
-    def _apply_to_column_groupbys(self, func, obj: FrameOrSeries) -> DataFrame:
+    def _apply_to_column_groupbys(self, func, obj: DataFrame | Series) -> DataFrame:
         from pandas.core.reshape.concat import concat
 
         columns = obj.columns

pandas/core/groupby/groupby.py

@@ -44,7 +44,6 @@ class providing the base-class of operations.
 import pandas._libs.groupby as libgroupby
 from pandas._typing import (
     ArrayLike,
-    F,
     FrameOrSeries,
     IndexLabel,
     RandomState,
@@ -1373,7 +1372,10 @@ def f(g):
 
     @final
     def _python_apply_general(
-        self, f: F, data: DataFrame | Series, not_indexed_same: bool | None = None
+        self,
+        f: Callable,
+        data: DataFrame | Series,
+        not_indexed_same: bool | None = None,
     ) -> DataFrame | Series:
         """
         Apply function f in python space

pandas/core/groupby/grouper.py

@@ -5,6 +5,8 @@
 from __future__ import annotations
 
 from typing import (
+    TYPE_CHECKING,
+    Any,
     Hashable,
     final,
 )
@@ -47,6 +49,9 @@
 
 from pandas.io.formats.printing import pprint_thing
 
+if TYPE_CHECKING:
+    from pandas.core.generic import NDFrame
+
 
 class Grouper:
     """
@@ -299,7 +304,9 @@ def ax(self) -> Index:
             raise ValueError("_set_grouper must be called before ax is accessed")
         return index
 
-    def _get_grouper(self, obj: FrameOrSeries, validate: bool = True):
+    def _get_grouper(
+        self, obj: FrameOrSeries, validate: bool = True
+    ) -> tuple[Any, ops.BaseGrouper, FrameOrSeries]:
         """
         Parameters
         ----------
@@ -326,10 +333,12 @@ def _get_grouper(self, obj: FrameOrSeries, validate: bool = True):
             dropna=self.dropna,
         )
 
-        return self.binner, self.grouper, self.obj
+        # error: Incompatible return value type (got "Tuple[None, None, None]",
+        # expected "Tuple[Any, BaseGrouper, FrameOrSeries]")
+        return self.binner, self.grouper, self.obj  # type: ignore[return-value]
 
     @final
-    def _set_grouper(self, obj: FrameOrSeries, sort: bool = False):
+    def _set_grouper(self, obj: NDFrame, sort: bool = False):
         """
         given an object and the specifications, setup the internal grouper
         for this particular specification
@@ -459,7 +468,7 @@ def __init__(
         self,
         index: Index,
         grouper=None,
-        obj: FrameOrSeries | None = None,
+        obj: NDFrame | None = None,
         level=None,
         sort: bool = True,
         observed: bool = False,
@@ -501,11 +510,9 @@ def __init__(
             # what key/level refer to exactly, don't need to
             # check again as we have by this point converted these
             # to an actual value (rather than a pd.Grouper)
+            assert self.obj is not None  # for mypy
             _, newgrouper, newobj = self.grouping_vector._get_grouper(
-                # error: Value of type variable "FrameOrSeries" of "_get_grouper"
-                # of "Grouper" cannot be "Optional[FrameOrSeries]"
-                self.obj,  # type: ignore[type-var]
-                validate=False,
+                self.obj, validate=False
             )
             self.obj = newobj
 

pandas/core/groupby/ops.py

@@ -10,6 +10,7 @@
 import collections
 import functools
 from typing import (
+    Callable,
     Generic,
     Hashable,
     Iterator,
@@ -29,7 +30,6 @@
 from pandas._typing import (
     ArrayLike,
     DtypeObj,
-    F,
     FrameOrSeries,
     Shape,
     npt,
@@ -700,7 +700,7 @@ def get_iterator(
             yield key, group.__finalize__(data, method="groupby")
 
     @final
-    def _get_splitter(self, data: FrameOrSeries, axis: int = 0) -> DataSplitter:
+    def _get_splitter(self, data: NDFrame, axis: int = 0) -> DataSplitter:
         """
         Returns
         -------
@@ -732,7 +732,9 @@ def group_keys_seq(self):
             return get_flattened_list(ids, ngroups, self.levels, self.codes)
 
     @final
-    def apply(self, f: F, data: FrameOrSeries, axis: int = 0) -> tuple[list, bool]:
+    def apply(
+        self, f: Callable, data: DataFrame | Series, axis: int = 0
+    ) -> tuple[list, bool]:
         mutated = self.mutated
         splitter = self._get_splitter(data, axis=axis)
         group_keys = self.group_keys_seq
@@ -918,7 +920,7 @@ def _cython_operation(
 
     @final
     def agg_series(
-        self, obj: Series, func: F, preserve_dtype: bool = False
+        self, obj: Series, func: Callable, preserve_dtype: bool = False
     ) -> ArrayLike:
         """
         Parameters
@@ -960,7 +962,7 @@ def agg_series(
 
     @final
     def _aggregate_series_pure_python(
-        self, obj: Series, func: F
+        self, obj: Series, func: Callable
     ) -> npt.NDArray[np.object_]:
         ids, _, ngroups = self.group_info
 
@@ -1061,7 +1063,7 @@ def _get_grouper(self):
         """
         return self
 
-    def get_iterator(self, data: FrameOrSeries, axis: int = 0):
+    def get_iterator(self, data: NDFrame, axis: int = 0):
         """
         Groupby iterator
 
@@ -1142,7 +1144,7 @@ def groupings(self) -> list[grouper.Grouping]:
         ping = grouper.Grouping(lev, lev, in_axis=False, level=None)
         return [ping]
 
-    def _aggregate_series_fast(self, obj: Series, func: F) -> np.ndarray:
+    def _aggregate_series_fast(self, obj: Series, func: Callable) -> np.ndarray:
         # -> np.ndarray[object]
         raise NotImplementedError(
             "This should not be reached; use _aggregate_series_pure_python"
@@ -1241,7 +1243,7 @@ def _chop(self, sdata: DataFrame, slice_obj: slice) -> DataFrame:
 
 
 def get_splitter(
-    data: FrameOrSeries, labels: np.ndarray, ngroups: int, axis: int = 0
+    data: NDFrame, labels: np.ndarray, ngroups: int, axis: int = 0
 ) -> DataSplitter:
     if isinstance(data, Series):
         klass: type[DataSplitter] = SeriesSplitter

pandas/core/resample.py

@@ -4,6 +4,7 @@
 from datetime import timedelta
 from textwrap import dedent
 from typing import (
+    TYPE_CHECKING,
     Callable,
     Hashable,
     Literal,
@@ -90,6 +91,12 @@
     Tick,
 )
 
+if TYPE_CHECKING:
+    from pandas import (
+        DataFrame,
+        Series,
+    )
+
 _shared_docs_kwargs: dict[str, str] = {}
 
 
@@ -135,7 +142,7 @@ class Resampler(BaseGroupBy, PandasObject):
 
     def __init__(
         self,
-        obj: FrameOrSeries,
+        obj: DataFrame | Series,
         groupby: TimeGrouper,
         axis: int = 0,
         kind=None,

pandas/core/reshape/merge.py

@@ -26,7 +26,6 @@
 from pandas._typing import (
     ArrayLike,
     DtypeObj,
-    FrameOrSeries,
     IndexLabel,
     Suffixes,
     npt,
@@ -2259,7 +2258,7 @@ def _any(x) -> bool:
     return x is not None and com.any_not_none(*x)
 
 
-def _validate_operand(obj: FrameOrSeries) -> DataFrame:
+def _validate_operand(obj: DataFrame | Series) -> DataFrame:
     if isinstance(obj, ABCDataFrame):
         return obj
     elif isinstance(obj, ABCSeries):

pandas/core/sample.py

@@ -3,18 +3,22 @@
 """
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import numpy as np
 
 from pandas._libs import lib
-from pandas._typing import FrameOrSeries
 
 from pandas.core.dtypes.generic import (
     ABCDataFrame,
     ABCSeries,
 )
 
+if TYPE_CHECKING:
+    from pandas.core.generic import NDFrame
+
 
-def preprocess_weights(obj: FrameOrSeries, weights, axis: int) -> np.ndarray:
+def preprocess_weights(obj: NDFrame, weights, axis: int) -> np.ndarray:
     """
     Process and validate the `weights` argument to `NDFrame.sample` and
     `.GroupBy.sample`.