TYP: Check for use of Union[Series, DataFrame] instead of FrameOrSeriesUnion alias (#36137)

Author: simonjayhawkins

ci/code_checks.sh

@@ -230,6 +230,9 @@ if [[ -z "$CHECK" || "$CHECK" == "patterns" ]]; then
     invgrep -R --include=*.{py,pyx} '!r}' pandas
     RET=$(($RET + $?)) ; echo $MSG "DONE"
 
+    # -------------------------------------------------------------------------
+    # Type annotations
+
     MSG='Check for use of comment-based annotation syntax' ; echo $MSG
     invgrep -R --include="*.py" -P '# type: (?!ignore)' pandas
     RET=$(($RET + $?)) ; echo $MSG "DONE"
@@ -238,6 +241,11 @@ if [[ -z "$CHECK" || "$CHECK" == "patterns" ]]; then
     invgrep -R --include="*.py" -P '# type:\s?ignore(?!\[)' pandas
     RET=$(($RET + $?)) ; echo $MSG "DONE"
 
+    MSG='Check for use of Union[Series, DataFrame] instead of FrameOrSeriesUnion alias' ; echo $MSG
+    invgrep -R --include="*.py" --exclude=_typing.py -E 'Union\[.*(Series.*DataFrame|DataFrame.*Series).*\]' pandas
+    RET=$(($RET + $?)) ; echo $MSG "DONE"
+
+    # -------------------------------------------------------------------------
     MSG='Check for use of foo.__class__ instead of type(foo)' ; echo $MSG
     invgrep -R --include=*.{py,pyx} '\.__class__' pandas
     RET=$(($RET + $?)) ; echo $MSG "DONE"

pandas/core/apply.py

@@ -1,12 +1,12 @@
 import abc
 import inspect
-from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional, Tuple, Type, Union
+from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional, Tuple, Type
 
 import numpy as np
 
 from pandas._config import option_context
 
-from pandas._typing import Axis
+from pandas._typing import Axis, FrameOrSeriesUnion
 from pandas.util._decorators import cache_readonly
 
 from pandas.core.dtypes.common import is_dict_like, is_list_like, is_sequence
@@ -73,7 +73,7 @@ def series_generator(self) -> Iterator["Series"]:
     @abc.abstractmethod
     def wrap_results_for_axis(
         self, results: ResType, res_index: "Index"
-    ) -> Union["Series", "DataFrame"]:
+    ) -> FrameOrSeriesUnion:
         pass
 
     # ---------------------------------------------------------------
@@ -289,9 +289,7 @@ def apply_series_generator(self) -> Tuple[ResType, "Index"]:
 
         return results, res_index
 
-    def wrap_results(
-        self, results: ResType, res_index: "Index"
-    ) -> Union["Series", "DataFrame"]:
+    def wrap_results(self, results: ResType, res_index: "Index") -> FrameOrSeriesUnion:
         from pandas import Series
 
         # see if we can infer the results
@@ -335,7 +333,7 @@ def result_columns(self) -> "Index":
 
     def wrap_results_for_axis(
         self, results: ResType, res_index: "Index"
-    ) -> Union["Series", "DataFrame"]:
+    ) -> FrameOrSeriesUnion:
         """ return the results for the rows """
 
         if self.result_type == "reduce":
@@ -408,9 +406,9 @@ def result_columns(self) -> "Index":
 
     def wrap_results_for_axis(
         self, results: ResType, res_index: "Index"
-    ) -> Union["Series", "DataFrame"]:
+    ) -> FrameOrSeriesUnion:
         """ return the results for the columns """
-        result: Union["Series", "DataFrame"]
+        result: FrameOrSeriesUnion
 
         # we have requested to expand
         if self.result_type == "expand":

pandas/core/groupby/generic.py

@@ -308,7 +308,7 @@ def _aggregate_multiple_funcs(self, arg):
 
             arg = zip(columns, arg)
 
-        results: Dict[base.OutputKey, Union[Series, DataFrame]] = {}
+        results: Dict[base.OutputKey, FrameOrSeriesUnion] = {}
         for idx, (name, func) in enumerate(arg):
             obj = self
 
@@ -332,7 +332,7 @@ def _wrap_series_output(
         self,
         output: Mapping[base.OutputKey, Union[Series, np.ndarray]],
         index: Optional[Index],
-    ) -> Union[Series, DataFrame]:
+    ) -> FrameOrSeriesUnion:
         """
         Wraps the output of a SeriesGroupBy operation into the expected result.
 
@@ -355,7 +355,7 @@ def _wrap_series_output(
         indexed_output = {key.position: val for key, val in output.items()}
         columns = Index(key.label for key in output)
 
-        result: Union[Series, DataFrame]
+        result: FrameOrSeriesUnion
         if len(output) > 1:
             result = self.obj._constructor_expanddim(indexed_output, index=index)
             result.columns = columns
@@ -373,7 +373,7 @@ def _wrap_aggregated_output(
         self,
         output: Mapping[base.OutputKey, Union[Series, np.ndarray]],
         index: Optional[Index],
-    ) -> Union[Series, DataFrame]:
+    ) -> FrameOrSeriesUnion:
         """
         Wraps the output of a SeriesGroupBy aggregation into the expected result.
 
@@ -1085,7 +1085,7 @@ def blk_func(bvalues: ArrayLike) -> ArrayLike:
                     raise
 
                 # We get here with a) EADtypes and b) object dtype
-                obj: Union[Series, DataFrame]
+                obj: FrameOrSeriesUnion
                 # call our grouper again with only this block
                 if isinstance(bvalues, ExtensionArray):
                     # TODO(EA2D): special case not needed with 2D EAs

pandas/core/groupby/grouper.py

@@ -393,7 +393,7 @@ class Grouping:
     ----------
     index : Index
     grouper :
-    obj Union[DataFrame, Series]:
+    obj : DataFrame or Series
     name : Label
     level :
     observed : bool, default False

pandas/core/reshape/merge.py

@@ -6,14 +6,14 @@
 import datetime
 from functools import partial
 import string
-from typing import TYPE_CHECKING, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Optional, Tuple
 import warnings
 
 import numpy as np
 
 from pandas._libs import Timedelta, hashtable as libhashtable, lib
 import pandas._libs.join as libjoin
-from pandas._typing import ArrayLike, FrameOrSeries
+from pandas._typing import ArrayLike, FrameOrSeries, FrameOrSeriesUnion
 from pandas.errors import MergeError
 from pandas.util._decorators import Appender, Substitution
 
@@ -51,7 +51,7 @@
 from pandas.core.sorting import is_int64_overflow_possible
 
 if TYPE_CHECKING:
-    from pandas import DataFrame, Series  # noqa:F401
+    from pandas import DataFrame  # noqa:F401
 
 
 @Substitution("\nleft : DataFrame")
@@ -575,8 +575,8 @@ class _MergeOperation:
 
     def __init__(
         self,
-        left: Union["Series", "DataFrame"],
-        right: Union["Series", "DataFrame"],
+        left: FrameOrSeriesUnion,
+        right: FrameOrSeriesUnion,
         how: str = "inner",
         on=None,
         left_on=None,

pandas/core/reshape/pivot.py

@@ -12,7 +12,7 @@
 
 import numpy as np
 
-from pandas._typing import Label
+from pandas._typing import FrameOrSeriesUnion, Label
 from pandas.util._decorators import Appender, Substitution
 
 from pandas.core.dtypes.cast import maybe_downcast_to_dtype
@@ -200,7 +200,7 @@ def pivot_table(
 
 
 def _add_margins(
-    table: Union["Series", "DataFrame"],
+    table: FrameOrSeriesUnion,
     data,
     values,
     rows,

pandas/io/pytables.py

@@ -16,7 +16,7 @@
 
 from pandas._libs import lib, writers as libwriters
 from pandas._libs.tslibs import timezones
-from pandas._typing import ArrayLike, FrameOrSeries, Label
+from pandas._typing import ArrayLike, FrameOrSeries, FrameOrSeriesUnion, Label
 from pandas.compat._optional import import_optional_dependency
 from pandas.compat.pickle_compat import patch_pickle
 from pandas.errors import PerformanceWarning
@@ -2566,7 +2566,7 @@ class Fixed:
 
     pandas_kind: str
     format_type: str = "fixed"  # GH#30962 needed by dask
-    obj_type: Type[Union[DataFrame, Series]]
+    obj_type: Type[FrameOrSeriesUnion]
     ndim: int
     encoding: str
     parent: HDFStore
@@ -4442,7 +4442,7 @@ class AppendableFrameTable(AppendableTable):
     pandas_kind = "frame_table"
     table_type = "appendable_frame"
     ndim = 2
-    obj_type: Type[Union[DataFrame, Series]] = DataFrame
+    obj_type: Type[FrameOrSeriesUnion] = DataFrame
 
     @property
     def is_transposed(self) -> bool: