Explain typed etc.

Author: topper-123

pandas/_typing.py

@@ -42,8 +42,16 @@
 
 Dtype = Union[str, np.dtype, "ExtensionDtype"]
 FilePathOrBuffer = Union[str, Path, IO[AnyStr]]
+# FrameOrSeries means either a DataFrame or a Series. E.g.
+# `def func(a: FrameOrSeries) -> FrameOrSeries: ...` means that if a Series is passed
+# in, either a Series or DataFrame is returned, and if a DataFrame is passed in, either
+# a DataFrame or a Series is returned.
 FrameOrSeries = Union["DataFrame", "Series"]
-FrameOrSeriesT = TypeVar("FrameOrSeriesT", bound="NDFrame")
+# SameFrameOrSeries is stricter and ensures that the same subclass of NDFrame always is
+# used. E.g. `def func(a: SameFrameOrSeries) -> SameFrameOrSeries: ...` means that if a
+# Series is passed into a function, a Series is always returned and if a DataFrame is
+# passed in, a DataFrame is always returned.
+SameFrameOrSeries = TypeVar("SameFrameOrSeries", bound="NDFrame")
 Axis = Union[str, int]
 Ordered = Optional[bool]
 JSONSerializable = Union[PythonScalar, List, Dict]

pandas/core/generic.py

@@ -29,7 +29,7 @@
 from pandas._config import config
 
 from pandas._libs import Timestamp, iNaT, lib, properties
-from pandas._typing import Dtype, FilePathOrBuffer, FrameOrSeriesT, JSONSerializable
+from pandas._typing import Dtype, FilePathOrBuffer, JSONSerializable, SameFrameOrSeries
 from pandas.compat import set_function_name
 from pandas.compat._optional import import_optional_dependency
 from pandas.compat.numpy import function as nv
@@ -542,12 +542,12 @@ def size(self):
         return np.prod(self.shape)
 
     @property
-    def _selected_obj(self: FrameOrSeriesT) -> FrameOrSeriesT:
+    def _selected_obj(self: SameFrameOrSeries) -> SameFrameOrSeries:
         """ internal compat with SelectionMixin """
         return self
 
     @property
-    def _obj_with_exclusions(self: FrameOrSeriesT) -> FrameOrSeriesT:
+    def _obj_with_exclusions(self: SameFrameOrSeries) -> SameFrameOrSeries:
         """ internal compat with SelectionMixin """
         return self
 
@@ -4644,7 +4644,7 @@ def f(x):
         else:
             raise TypeError("Must pass either `items`, `like`, or `regex`")
 
-    def head(self: FrameOrSeriesT, n: int = 5) -> FrameOrSeriesT:
+    def head(self: SameFrameOrSeries, n: int = 5) -> SameFrameOrSeries:
         """
         Return the first `n` rows.
 
@@ -4717,7 +4717,7 @@ def head(self: FrameOrSeriesT, n: int = 5) -> FrameOrSeriesT:
 
         return self.iloc[:n]
 
-    def tail(self: FrameOrSeriesT, n: int = 5) -> FrameOrSeriesT:
+    def tail(self: SameFrameOrSeries, n: int = 5) -> SameFrameOrSeries:
         """
         Return the last `n` rows.
 
@@ -5162,8 +5162,8 @@ def pipe(self, func, *args, **kwargs):
     # Attribute access
 
     def __finalize__(
-        self: FrameOrSeriesT, other, method=None, **kwargs
-    ) -> FrameOrSeriesT:
+        self: SameFrameOrSeries, other, method=None, **kwargs
+    ) -> SameFrameOrSeries:
         """
         Propagate metadata from other to self.
 
@@ -5630,7 +5630,7 @@ def astype(self, dtype, copy: bool_t = True, errors: str = "raise"):
         result.columns = self.columns
         return result
 
-    def copy(self: FrameOrSeriesT, deep: bool_t = True) -> FrameOrSeriesT:
+    def copy(self: SameFrameOrSeries, deep: bool_t = True) -> SameFrameOrSeries:
         """
         Make a copy of this object's indices and data.
 
@@ -5738,10 +5738,10 @@ def copy(self: FrameOrSeriesT, deep: bool_t = True) -> FrameOrSeriesT:
         data = self._data.copy(deep=deep)
         return self._constructor(data).__finalize__(self)
 
-    def __copy__(self: FrameOrSeriesT, deep: bool_t = True) -> FrameOrSeriesT:
+    def __copy__(self: SameFrameOrSeries, deep: bool_t = True) -> SameFrameOrSeries:
         return self.copy(deep=deep)
 
-    def __deepcopy__(self: FrameOrSeriesT, memo=None) -> FrameOrSeriesT:
+    def __deepcopy__(self: SameFrameOrSeries, memo=None) -> SameFrameOrSeries:
         """
         Parameters
         ----------
@@ -5751,13 +5751,13 @@ def __deepcopy__(self: FrameOrSeriesT, memo=None) -> FrameOrSeriesT:
         return self.copy(deep=True)
 
     def _convert(
-        self: FrameOrSeriesT,
+        self: SameFrameOrSeries,
         datetime: bool_t = False,
         numeric: bool_t = False,
         timedelta: bool_t = False,
         coerce: bool_t = False,
         copy: bool_t = True,
-    ) -> FrameOrSeriesT:
+    ) -> SameFrameOrSeries:
         """
         Attempt to infer better dtype for object columns
 
@@ -5849,14 +5849,14 @@ def infer_objects(self):
     # Filling NA's
 
     def fillna(
-        self: FrameOrSeriesT,
+        self: SameFrameOrSeries,
         value=None,
         method=None,
         axis=None,
         inplace: bool_t = False,
         limit=None,
         downcast=None,
-    ) -> Optional[FrameOrSeriesT]:
+    ) -> Optional[SameFrameOrSeries]:
         """
         Fill NA/NaN values using the specified method.
 
@@ -6038,12 +6038,12 @@ def fillna(
             return self._constructor(new_data).__finalize__(self)
 
     def ffill(
-        self: FrameOrSeriesT,
+        self: SameFrameOrSeries,
         axis=None,
         inplace: bool_t = False,
         limit=None,
         downcast=None,
-    ) -> Optional[FrameOrSeriesT]:
+    ) -> Optional[SameFrameOrSeries]:
         """
         Synonym for :meth:`DataFrame.fillna` with ``method='ffill'``.
 
@@ -6057,12 +6057,12 @@ def ffill(
         )
 
     def bfill(
-        self: FrameOrSeriesT,
+        self: SameFrameOrSeries,
         axis=None,
         inplace: bool_t = False,
         limit=None,
         downcast=None,
-    ) -> Optional[FrameOrSeriesT]:
+    ) -> Optional[SameFrameOrSeries]:
         """
         Synonym for :meth:`DataFrame.fillna` with ``method='bfill'``.
 
@@ -8025,14 +8025,14 @@ def last(self, offset):
         return self.iloc[start:]
 
     def rank(
-        self: FrameOrSeriesT,
+        self: SameFrameOrSeries,
         axis=0,
         method: str = "average",
         numeric_only: Optional[bool_t] = None,
         na_option: str = "keep",
         ascending: bool_t = True,
         pct: bool_t = False,
-    ) -> FrameOrSeriesT:
+    ) -> SameFrameOrSeries:
         """
         Compute numerical data ranks (1 through n) along axis.
 
@@ -8834,7 +8834,9 @@ def shift(self, periods=1, freq=None, axis=0, fill_value=None):
 
         return self._constructor(new_data).__finalize__(self)
 
-    def slice_shift(self: FrameOrSeriesT, periods: int = 1, axis=0) -> FrameOrSeriesT:
+    def slice_shift(
+        self: SameFrameOrSeries, periods: int = 1, axis=0
+    ) -> SameFrameOrSeries:
         """
         Equivalent to `shift` without copying data.
 
@@ -8932,8 +8934,8 @@ def tshift(self, periods: int = 1, freq=None, axis=0):
         return self._constructor(new_data).__finalize__(self)
 
     def truncate(
-        self: FrameOrSeriesT, before=None, after=None, axis=None, copy: bool_t = True
-    ) -> FrameOrSeriesT:
+        self: SameFrameOrSeries, before=None, after=None, axis=None, copy: bool_t = True
+    ) -> SameFrameOrSeries:
         """
         Truncate a Series or DataFrame before and after some index value.
 
@@ -9086,8 +9088,8 @@ def truncate(
         return result
 
     def tz_convert(
-        self: FrameOrSeriesT, tz, axis=0, level=None, copy: bool_t = True
-    ) -> FrameOrSeriesT:
+        self: SameFrameOrSeries, tz, axis=0, level=None, copy: bool_t = True
+    ) -> SameFrameOrSeries:
         """
         Convert tz-aware axis to target time zone.
 
@@ -9143,14 +9145,14 @@ def _tz_convert(ax, tz):
         return result.__finalize__(self)
 
     def tz_localize(
-        self: FrameOrSeriesT,
+        self: SameFrameOrSeries,
         tz,
         axis=0,
         level=None,
         copy: bool_t = True,
         ambiguous="raise",
         nonexistent: str = "raise",
-    ) -> FrameOrSeriesT:
+    ) -> SameFrameOrSeries:
         """
         Localize tz-naive index of a Series or DataFrame to target time zone.
 

pandas/core/groupby/generic.py

@@ -29,7 +29,7 @@
 import numpy as np
 
 from pandas._libs import Timestamp, lib
-from pandas._typing import FrameOrSeriesT
+from pandas._typing import SameFrameOrSeries
 from pandas.util._decorators import Appender, Substitution
 
 from pandas.core.dtypes.cast import (
@@ -85,7 +85,7 @@
 ScalarResult = typing.TypeVar("ScalarResult")
 
 
-def generate_property(name: str, klass: Type[FrameOrSeriesT]):
+def generate_property(name: str, klass: Type[SameFrameOrSeries]):
     """
     Create a property for a GroupBy subclass to dispatch to DataFrame/Series.
 
@@ -108,7 +108,9 @@ def prop(self):
     return property(prop)
 
 
-def pin_whitelisted_properties(klass: Type[FrameOrSeriesT], whitelist: FrozenSet[str]):
+def pin_whitelisted_properties(
+    klass: Type[SameFrameOrSeries], whitelist: FrozenSet[str]
+):
     """
     Create GroupBy member defs for DataFrame/Series names in a whitelist.
 

pandas/core/groupby/groupby.py

@@ -33,7 +33,7 @@ class providing the base-class of operations.
 
 from pandas._libs import Timestamp
 import pandas._libs.groupby as libgroupby
-from pandas._typing import FrameOrSeriesT, Scalar
+from pandas._typing import SameFrameOrSeries, Scalar
 from pandas.compat import set_function_name
 from pandas.compat.numpy import function as nv
 from pandas.errors import AbstractMethodError
@@ -2439,8 +2439,8 @@ def tail(self, n=5):
         return self._selected_obj[mask]
 
     def _reindex_output(
-        self, output: FrameOrSeriesT, fill_value: Scalar = np.NaN
-    ) -> FrameOrSeriesT:
+        self, output: SameFrameOrSeries, fill_value: Scalar = np.NaN
+    ) -> SameFrameOrSeries:
         """
         If we have categorical groupers, then we might want to make sure that
         we have a fully re-indexed output to the levels. This means expanding

pandas/core/groupby/grouper.py

@@ -7,7 +7,7 @@
 
 import numpy as np
 
-from pandas._typing import FrameOrSeriesT
+from pandas._typing import SameFrameOrSeries
 from pandas.util._decorators import cache_readonly
 
 from pandas.core.dtypes.common import (
@@ -141,7 +141,7 @@ def _get_grouper(self, obj, validate: bool = True):
         )
         return self.binner, self.grouper, self.obj
 
-    def _set_grouper(self, obj: FrameOrSeriesT, sort: bool = False):
+    def _set_grouper(self, obj: SameFrameOrSeries, sort: bool = False):
         """
         given an object and the specifications, setup the internal grouper
         for this particular specification
@@ -244,7 +244,7 @@ def __init__(
         self,
         index: Index,
         grouper=None,
-        obj: Optional[FrameOrSeriesT] = None,
+        obj: Optional[SameFrameOrSeries] = None,
         name=None,
         level=None,
         sort: bool = True,
@@ -424,15 +424,15 @@ def groups(self) -> Dict[Hashable, np.ndarray]:
 
 
 def get_grouper(
-    obj: FrameOrSeriesT,
+    obj: SameFrameOrSeries,
     key=None,
     axis: int = 0,
     level=None,
     sort: bool = True,
     observed: bool = False,
     mutated: bool = False,
     validate: bool = True,
-) -> "Tuple[ops.BaseGrouper, List[Hashable], FrameOrSeriesT]":
+) -> "Tuple[ops.BaseGrouper, List[Hashable], SameFrameOrSeries]":
     """
     Create and return a BaseGrouper, which is an internal
     mapping of how to create the grouper indexers.

pandas/core/groupby/ops.py

@@ -14,7 +14,7 @@
 from pandas._libs import NaT, iNaT, lib
 import pandas._libs.groupby as libgroupby
 import pandas._libs.reduction as libreduction
-from pandas._typing import FrameOrSeriesT
+from pandas._typing import SameFrameOrSeries
 from pandas.errors import AbstractMethodError
 from pandas.util._decorators import cache_readonly
 
@@ -111,7 +111,7 @@ def __iter__(self):
     def nkeys(self) -> int:
         return len(self.groupings)
 
-    def get_iterator(self, data: FrameOrSeriesT, axis: int = 0):
+    def get_iterator(self, data: SameFrameOrSeries, axis: int = 0):
         """
         Groupby iterator
 
@@ -125,7 +125,7 @@ def get_iterator(self, data: FrameOrSeriesT, axis: int = 0):
         for key, (i, group) in zip(keys, splitter):
             yield key, group
 
-    def _get_splitter(self, data: FrameOrSeriesT, axis: int = 0) -> "DataSplitter":
+    def _get_splitter(self, data: SameFrameOrSeries, axis: int = 0) -> "DataSplitter":
         comp_ids, _, ngroups = self.group_info
         return get_splitter(data, comp_ids, ngroups, axis=axis)
 
@@ -147,13 +147,13 @@ def _get_group_keys(self):
             # provide "flattened" iterator for multi-group setting
             return get_flattened_iterator(comp_ids, ngroups, self.levels, self.codes)
 
-    def apply(self, f, data: FrameOrSeriesT, axis: int = 0):
+    def apply(self, f, data: SameFrameOrSeries, axis: int = 0):
         mutated = self.mutated
         splitter = self._get_splitter(data, axis=axis)
         group_keys = self._get_group_keys()
         result_values = None
 
-        sdata: FrameOrSeriesT = splitter._get_sorted_data()
+        sdata: SameFrameOrSeries = splitter._get_sorted_data()
         if sdata.ndim == 2 and np.any(sdata.dtypes.apply(is_extension_array_dtype)):
             # calling splitter.fast_apply will raise TypeError via apply_frame_axis0
             #  if we pass EA instead of ndarray
@@ -754,7 +754,7 @@ def _get_grouper(self):
         """
         return self
 
-    def get_iterator(self, data: FrameOrSeriesT, axis: int = 0):
+    def get_iterator(self, data: SameFrameOrSeries, axis: int = 0):
         """
         Groupby iterator
 
@@ -862,7 +862,7 @@ def _is_indexed_like(obj, axes) -> bool:
 
 
 class DataSplitter:
-    def __init__(self, data: FrameOrSeriesT, labels, ngroups: int, axis: int = 0):
+    def __init__(self, data: SameFrameOrSeries, labels, ngroups: int, axis: int = 0):
         self.data = data
         self.labels = ensure_int64(labels)
         self.ngroups = ngroups
@@ -893,7 +893,7 @@ def __iter__(self):
         for i, (start, end) in enumerate(zip(starts, ends)):
             yield i, self._chop(sdata, slice(start, end))
 
-    def _get_sorted_data(self) -> FrameOrSeriesT:
+    def _get_sorted_data(self) -> SameFrameOrSeries:
         return self.data.take(self.sort_idx, axis=self.axis)
 
     def _chop(self, sdata, slice_obj: slice) -> NDFrame:
@@ -920,7 +920,7 @@ def _chop(self, sdata: DataFrame, slice_obj: slice) -> DataFrame:
             return sdata._slice(slice_obj, axis=1)
 
 
-def get_splitter(data: FrameOrSeriesT, *args, **kwargs) -> DataSplitter:
+def get_splitter(data: SameFrameOrSeries, *args, **kwargs) -> DataSplitter:
     if isinstance(data, Series):
         klass: Type[DataSplitter] = SeriesSplitter
     else: