CLN: Replace FrameOrSeriesUnion annotation by DataFrame | Series (#41955)

Author: Swanand01

pandas/_typing.py

@@ -104,12 +104,6 @@
 ]
 Timezone = Union[str, tzinfo]
 
-# FrameOrSeriesUnion  means either a DataFrame or a Series. E.g.
-# `def func(a: FrameOrSeriesUnion) -> FrameOrSeriesUnion: ...` 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.
-FrameOrSeriesUnion = Union["DataFrame", "Series"]
-
 # FrameOrSeries is stricter and ensures that the same subclass of NDFrame always is
 # used. E.g. `def func(a: FrameOrSeries) -> FrameOrSeries: ...` means that if a
 # Series is passed into a function, a Series is always returned and if a DataFrame is

pandas/core/algorithms.py

@@ -25,7 +25,6 @@
     AnyArrayLike,
     ArrayLike,
     DtypeObj,
-    FrameOrSeriesUnion,
     Scalar,
 )
 from pandas.util._decorators import doc
@@ -1211,7 +1210,7 @@ def __init__(self, obj, n: int, keep: str):
         if self.keep not in ("first", "last", "all"):
             raise ValueError('keep must be either "first", "last" or "all"')
 
-    def compute(self, method: str) -> FrameOrSeriesUnion:
+    def compute(self, method: str) -> DataFrame | Series:
         raise NotImplementedError
 
     def nlargest(self):

pandas/core/apply.py

@@ -25,7 +25,6 @@
     AggObjType,
     Axis,
     FrameOrSeries,
-    FrameOrSeriesUnion,
 )
 from pandas.util._decorators import cache_readonly
 
@@ -137,10 +136,10 @@ def f(x):
         self.f: AggFuncType = f
 
     @abc.abstractmethod
-    def apply(self) -> FrameOrSeriesUnion:
+    def apply(self) -> DataFrame | Series:
         pass
 
-    def agg(self) -> FrameOrSeriesUnion | None:
+    def agg(self) -> DataFrame | Series | None:
         """
         Provide an implementation for the aggregators.
 
@@ -171,7 +170,7 @@ def agg(self) -> FrameOrSeriesUnion | None:
         # caller can react
         return None
 
-    def transform(self) -> FrameOrSeriesUnion:
+    def transform(self) -> DataFrame | Series:
         """
         Transform a DataFrame or Series.
 
@@ -252,7 +251,7 @@ def transform_dict_like(self, func):
 
         func = self.normalize_dictlike_arg("transform", obj, func)
 
-        results: dict[Hashable, FrameOrSeriesUnion] = {}
+        results: dict[Hashable, DataFrame | Series] = {}
         failed_names = []
         all_type_errors = True
         for name, how in func.items():
@@ -283,7 +282,7 @@ def transform_dict_like(self, func):
             )
         return concat(results, axis=1)
 
-    def transform_str_or_callable(self, func) -> FrameOrSeriesUnion:
+    def transform_str_or_callable(self, func) -> DataFrame | Series:
         """
         Compute transform in the case of a string or callable func
         """
@@ -305,7 +304,7 @@ def transform_str_or_callable(self, func) -> FrameOrSeriesUnion:
         except Exception:
             return func(obj, *args, **kwargs)
 
-    def agg_list_like(self) -> FrameOrSeriesUnion:
+    def agg_list_like(self) -> DataFrame | Series:
         """
         Compute aggregation in the case of a list-like argument.
 
@@ -402,7 +401,7 @@ def agg_list_like(self) -> FrameOrSeriesUnion:
             )
             return concatenated.reindex(full_ordered_index, copy=False)
 
-    def agg_dict_like(self) -> FrameOrSeriesUnion:
+    def agg_dict_like(self) -> DataFrame | Series:
         """
         Compute aggregation in the case of a dict-like argument.
 
@@ -481,7 +480,7 @@ def agg_dict_like(self) -> FrameOrSeriesUnion:
 
         return result
 
-    def apply_str(self) -> FrameOrSeriesUnion:
+    def apply_str(self) -> DataFrame | Series:
         """
         Compute apply in case of a string.
 
@@ -506,7 +505,7 @@ def apply_str(self) -> FrameOrSeriesUnion:
                 raise ValueError(f"Operation {f} does not support axis=1")
         return self._try_aggregate_string_function(obj, f, *self.args, **self.kwargs)
 
-    def apply_multiple(self) -> FrameOrSeriesUnion:
+    def apply_multiple(self) -> DataFrame | Series:
         """
         Compute apply in case of a list-like or dict-like.
 
@@ -518,7 +517,7 @@ def apply_multiple(self) -> FrameOrSeriesUnion:
         return self.obj.aggregate(self.f, self.axis, *self.args, **self.kwargs)
 
     def normalize_dictlike_arg(
-        self, how: str, obj: FrameOrSeriesUnion, func: AggFuncTypeDict
+        self, how: str, obj: DataFrame | Series, func: AggFuncTypeDict
     ) -> AggFuncTypeDict:
         """
         Handler for dict-like argument.
@@ -631,7 +630,7 @@ def series_generator(self) -> Iterator[Series]:
     @abc.abstractmethod
     def wrap_results_for_axis(
         self, results: ResType, res_index: Index
-    ) -> FrameOrSeriesUnion:
+    ) -> DataFrame | Series:
         pass
 
     # ---------------------------------------------------------------
@@ -652,7 +651,7 @@ def values(self):
     def dtypes(self) -> Series:
         return self.obj.dtypes
 
-    def apply(self) -> FrameOrSeriesUnion:
+    def apply(self) -> DataFrame | Series:
         """compute the results"""
         # dispatch to agg
         if is_list_like(self.f):
@@ -826,7 +825,7 @@ def apply_series_generator(self) -> tuple[ResType, Index]:
 
         return results, res_index
 
-    def wrap_results(self, results: ResType, res_index: Index) -> FrameOrSeriesUnion:
+    def wrap_results(self, results: ResType, res_index: Index) -> DataFrame | Series:
         from pandas import Series
 
         # see if we can infer the results
@@ -849,7 +848,7 @@ def wrap_results(self, results: ResType, res_index: Index) -> FrameOrSeriesUnion
 
         return result
 
-    def apply_str(self) -> FrameOrSeriesUnion:
+    def apply_str(self) -> DataFrame | Series:
         # Caller is responsible for checking isinstance(self.f, str)
         # TODO: GH#39993 - Avoid special-casing by replacing with lambda
         if self.f == "size":
@@ -880,7 +879,7 @@ def result_columns(self) -> Index:
 
     def wrap_results_for_axis(
         self, results: ResType, res_index: Index
-    ) -> FrameOrSeriesUnion:
+    ) -> DataFrame | Series:
         """return the results for the rows"""
 
         if self.result_type == "reduce":
@@ -963,9 +962,9 @@ def result_columns(self) -> Index:
 
     def wrap_results_for_axis(
         self, results: ResType, res_index: Index
-    ) -> FrameOrSeriesUnion:
+    ) -> DataFrame | Series:
         """return the results for the columns"""
-        result: FrameOrSeriesUnion
+        result: DataFrame | Series
 
         # we have requested to expand
         if self.result_type == "expand":
@@ -1019,7 +1018,7 @@ def __init__(
             kwargs=kwargs,
         )
 
-    def apply(self) -> FrameOrSeriesUnion:
+    def apply(self) -> DataFrame | Series:
         obj = self.obj
 
         if len(obj) == 0:
@@ -1070,7 +1069,7 @@ def apply_empty_result(self) -> Series:
             obj, method="apply"
         )
 
-    def apply_standard(self) -> FrameOrSeriesUnion:
+    def apply_standard(self) -> DataFrame | Series:
         f = self.f
         obj = self.obj
 

pandas/core/describe.py

@@ -22,7 +22,6 @@
 from pandas._libs.tslibs import Timestamp
 from pandas._typing import (
     FrameOrSeries,
-    FrameOrSeriesUnion,
     Hashable,
 )
 from pandas.util._validators import validate_percentile
@@ -107,12 +106,12 @@ class NDFrameDescriberAbstract(ABC):
         Whether to treat datetime dtypes as numeric.
     """
 
-    def __init__(self, obj: FrameOrSeriesUnion, datetime_is_numeric: bool):
+    def __init__(self, obj: DataFrame | Series, datetime_is_numeric: bool):
         self.obj = obj
         self.datetime_is_numeric = datetime_is_numeric
 
     @abstractmethod
-    def describe(self, percentiles: Sequence[float]) -> FrameOrSeriesUnion:
+    def describe(self, percentiles: Sequence[float]) -> DataFrame | Series:
         """Do describe either series or dataframe.
 
         Parameters

pandas/core/frame.py

@@ -58,7 +58,6 @@
     FillnaOptions,
     FloatFormatType,
     FormattersType,
-    FrameOrSeriesUnion,
     Frequency,
     IndexKeyFunc,
     IndexLabel,
@@ -1362,7 +1361,7 @@ def dot(self, other: Series) -> Series:
     def dot(self, other: DataFrame | Index | ArrayLike) -> DataFrame:
         ...
 
-    def dot(self, other: AnyArrayLike | FrameOrSeriesUnion) -> FrameOrSeriesUnion:
+    def dot(self, other: AnyArrayLike | DataFrame | Series) -> DataFrame | Series:
         """
         Compute the matrix multiplication between the DataFrame and other.
 
@@ -1478,13 +1477,13 @@ def __matmul__(self, other: Series) -> Series:
 
     @overload
     def __matmul__(
-        self, other: AnyArrayLike | FrameOrSeriesUnion
-    ) -> FrameOrSeriesUnion:
+        self, other: AnyArrayLike | DataFrame | Series
+    ) -> DataFrame | Series:
         ...
 
     def __matmul__(
-        self, other: AnyArrayLike | FrameOrSeriesUnion
-    ) -> FrameOrSeriesUnion:
+        self, other: AnyArrayLike | DataFrame | Series
+    ) -> DataFrame | Series:
         """
         Matrix multiplication using binary `@` operator in Python>=3.5.
         """
@@ -8432,8 +8431,8 @@ def _gotitem(
         self,
         key: IndexLabel,
         ndim: int,
-        subset: FrameOrSeriesUnion | None = None,
-    ) -> FrameOrSeriesUnion:
+        subset: DataFrame | Series | None = None,
+    ) -> DataFrame | Series:
         """
         Sub-classes to define. Return a sliced object.
 
@@ -8962,7 +8961,7 @@ def append(
 
     def join(
         self,
-        other: FrameOrSeriesUnion,
+        other: DataFrame | Series,
         on: IndexLabel | None = None,
         how: str = "left",
         lsuffix: str = "",
@@ -9092,7 +9091,7 @@ def join(
 
     def _join_compat(
         self,
-        other: FrameOrSeriesUnion,
+        other: DataFrame | Series,
         on: IndexLabel | None = None,
         how: str = "left",
         lsuffix: str = "",
@@ -9162,7 +9161,7 @@ def _join_compat(
     @Appender(_merge_doc, indents=2)
     def merge(
         self,
-        right: FrameOrSeriesUnion,
+        right: DataFrame | Series,
         how: str = "inner",
         on: IndexLabel | None = None,
         left_on: IndexLabel | None = None,
@@ -10759,7 +10758,7 @@ def _from_nested_dict(data) -> collections.defaultdict:
     return new_data
 
 
-def _reindex_for_setitem(value: FrameOrSeriesUnion, index: Index) -> ArrayLike:
+def _reindex_for_setitem(value: DataFrame | Series, index: Index) -> ArrayLike:
     # reindex if necessary
 
     if value.index.equals(index) or not len(index):

pandas/core/groupby/generic.py

@@ -33,7 +33,6 @@
 from pandas._typing import (
     ArrayLike,
     FrameOrSeries,
-    FrameOrSeriesUnion,
     Manager2D,
 )
 from pandas.util._decorators import (
@@ -296,7 +295,7 @@ def _aggregate_multiple_funcs(self, arg) -> DataFrame:
 
             arg = zip(columns, arg)
 
-        results: dict[base.OutputKey, FrameOrSeriesUnion] = {}
+        results: dict[base.OutputKey, DataFrame | Series] = {}
         for idx, (name, func) in enumerate(arg):
 
             key = base.OutputKey(label=name, position=idx)
@@ -422,7 +421,7 @@ def _wrap_applied_output(
         keys: Index,
         values: list[Any] | None,
         not_indexed_same: bool = False,
-    ) -> FrameOrSeriesUnion:
+    ) -> DataFrame | Series:
         """
         Wrap the output of SeriesGroupBy.apply into the expected result.
 
@@ -1193,7 +1192,7 @@ def _wrap_applied_output_series(
         not_indexed_same: bool,
         first_not_none,
         key_index,
-    ) -> FrameOrSeriesUnion:
+    ) -> DataFrame | Series:
         # this is to silence a DeprecationWarning
         # TODO: Remove when default dtype of empty Series is object
         kwargs = first_not_none._construct_axes_dict()

pandas/core/groupby/groupby.py

@@ -45,7 +45,6 @@ class providing the base-class of operations.
     ArrayLike,
     F,
     FrameOrSeries,
-    FrameOrSeriesUnion,
     IndexLabel,
     RandomState,
     Scalar,
@@ -730,7 +729,7 @@ def pipe(
     plot = property(GroupByPlot)
 
     @final
-    def get_group(self, name, obj=None) -> FrameOrSeriesUnion:
+    def get_group(self, name, obj=None) -> DataFrame | Series:
         """
         Construct DataFrame from group with provided name.
 
@@ -1269,8 +1268,8 @@ def f(g):
 
     @final
     def _python_apply_general(
-        self, f: F, data: FrameOrSeriesUnion
-    ) -> FrameOrSeriesUnion:
+        self, f: F, data: DataFrame | Series
+    ) -> DataFrame | Series:
         """
         Apply function f in python space
 
@@ -1792,7 +1791,7 @@ def sem(self, ddof: int = 1):
     @final
     @Substitution(name="groupby")
     @Appender(_common_see_also)
-    def size(self) -> FrameOrSeriesUnion:
+    def size(self) -> DataFrame | Series:
         """
         Compute group sizes.