TYP: Overload concat (pandas-dev#41184)

* TYP: Overload concat

* Fix NDFrame columns overload

* Fix FrameOrSeries

* Fix redundant casts

* Import Axis outside of TYPE_CHECKING

* Remove redundant overloads

* Mark mypy/issues/8354 return type issues

* Fix missing imports

* noop

Co-authored-by: Marco Gorelli <[email protected]>

Author: 2 people

pandas/core/arrays/categorical.py

@@ -2310,10 +2310,11 @@ def describe(self):
         counts = self.value_counts(dropna=False)
         freqs = counts / counts.sum()
 
+        from pandas import Index
         from pandas.core.reshape.concat import concat
 
         result = concat([counts, freqs], axis=1)
-        result.columns = ["counts", "freqs"]
+        result.columns = Index(["counts", "freqs"])
         result.index.name = "categories"
 
         return result

pandas/core/generic.py

@@ -5755,7 +5755,8 @@ def astype(
         # GH 19920: retain column metadata after concat
         result = concat(results, axis=1, copy=False)
         result.columns = self.columns
-        return result
+        # https://github.com/python/mypy/issues/8354
+        return cast(FrameOrSeries, result)
 
     @final
     def copy(self: FrameOrSeries, deep: bool_t = True) -> FrameOrSeries:
@@ -6118,7 +6119,8 @@ def convert_dtypes(
                 for col_name, col in self.items()
             ]
             if len(results) > 0:
-                return concat(results, axis=1, copy=False)
+                # https://github.com/python/mypy/issues/8354
+                return cast(FrameOrSeries, concat(results, axis=1, copy=False))
             else:
                 return self.copy()
 

pandas/core/groupby/generic.py

@@ -308,9 +308,7 @@ def _aggregate_multiple_funcs(self, arg) -> DataFrame:
             res_df = concat(
                 results.values(), axis=1, keys=[key.label for key in results.keys()]
             )
-            # error: Incompatible return value type (got "Union[DataFrame, Series]",
-            # expected "DataFrame")
-            return res_df  # type: ignore[return-value]
+            return res_df
 
         indexed_output = {key.position: val for key, val in results.items()}
         output = self.obj._constructor_expanddim(indexed_output, index=None)
@@ -547,9 +545,7 @@ def _transform_general(self, func: Callable, *args, **kwargs) -> Series:
             result = self.obj._constructor(dtype=np.float64)
 
         result.name = self.obj.name
-        # error: Incompatible return value type (got "Union[DataFrame, Series]",
-        # expected "Series")
-        return result  # type: ignore[return-value]
+        return result
 
     def _can_use_transform_fast(self, result) -> bool:
         return True

pandas/core/reshape/concat.py

@@ -8,13 +8,15 @@
     TYPE_CHECKING,
     Hashable,
     Iterable,
+    Literal,
     Mapping,
     cast,
     overload,
 )
 
 import numpy as np
 
+from pandas._typing import Axis
 from pandas.util._decorators import (
     cache_readonly,
     deprecate_nonkeyword_arguments,
@@ -57,25 +59,73 @@
 @overload
 def concat(
     objs: Iterable[DataFrame] | Mapping[Hashable, DataFrame],
-    axis=0,
-    join: str = "outer",
-    ignore_index: bool = False,
-    keys=None,
-    levels=None,
-    names=None,
-    verify_integrity: bool = False,
-    sort: bool = False,
-    copy: bool = True,
+    axis: Literal[0, "index"] = ...,
+    join: str = ...,
+    ignore_index: bool = ...,
+    keys=...,
+    levels=...,
+    names=...,
+    verify_integrity: bool = ...,
+    sort: bool = ...,
+    copy: bool = ...,
 ) -> DataFrame:
     ...
 
 
+@overload
+def concat(
+    objs: Iterable[Series] | Mapping[Hashable, Series],
+    axis: Literal[0, "index"] = ...,
+    join: str = ...,
+    ignore_index: bool = ...,
+    keys=...,
+    levels=...,
+    names=...,
+    verify_integrity: bool = ...,
+    sort: bool = ...,
+    copy: bool = ...,
+) -> Series:
+    ...
+
+
 @overload
 def concat(
     objs: Iterable[NDFrame] | Mapping[Hashable, NDFrame],
-    axis=0,
-    join: str = "outer",
-    ignore_index: bool = False,
+    axis: Literal[0, "index"] = ...,
+    join: str = ...,
+    ignore_index: bool = ...,
+    keys=...,
+    levels=...,
+    names=...,
+    verify_integrity: bool = ...,
+    sort: bool = ...,
+    copy: bool = ...,
+) -> DataFrame | Series:
+    ...
+
+
+@overload
+def concat(
+    objs: Iterable[NDFrame] | Mapping[Hashable, NDFrame],
+    axis: Literal[1, "columns"],
+    join: str = ...,
+    ignore_index: bool = ...,
+    keys=...,
+    levels=...,
+    names=...,
+    verify_integrity: bool = ...,
+    sort: bool = ...,
+    copy: bool = ...,
+) -> DataFrame:
+    ...
+
+
+@overload
+def concat(
+    objs: Iterable[NDFrame] | Mapping[Hashable, NDFrame],
+    axis: Axis = ...,
+    join: str = ...,
+    ignore_index: bool = ...,
     keys=None,
     levels=None,
     names=None,
@@ -89,8 +139,8 @@ def concat(
 @deprecate_nonkeyword_arguments(version=None, allowed_args=["objs"])
 def concat(
     objs: Iterable[NDFrame] | Mapping[Hashable, NDFrame],
-    axis=0,
-    join="outer",
+    axis: Axis = 0,
+    join: str = "outer",
     ignore_index: bool = False,
     keys=None,
     levels=None,

pandas/core/reshape/melt.py

@@ -1,10 +1,7 @@
 from __future__ import annotations
 
 import re
-from typing import (
-    TYPE_CHECKING,
-    cast,
-)
+from typing import TYPE_CHECKING
 import warnings
 
 import numpy as np
@@ -34,10 +31,7 @@
 from pandas.core.tools.numeric import to_numeric
 
 if TYPE_CHECKING:
-    from pandas import (
-        DataFrame,
-        Series,
-    )
+    from pandas import DataFrame
 
 
 @Appender(_shared_docs["melt"] % {"caller": "pd.melt(df, ", "other": "DataFrame.melt"})
@@ -136,7 +130,7 @@ def melt(
     for col in id_vars:
         id_data = frame.pop(col)
         if is_extension_array_dtype(id_data):
-            id_data = cast("Series", concat([id_data] * K, ignore_index=True))
+            id_data = concat([id_data] * K, ignore_index=True)
         else:
             id_data = np.tile(id_data._values, K)
         mdata[col] = id_data

pandas/core/reshape/reshape.py

@@ -1,10 +1,7 @@
 from __future__ import annotations
 
 import itertools
-from typing import (
-    TYPE_CHECKING,
-    cast,
-)
+from typing import TYPE_CHECKING
 
 import numpy as np
 
@@ -1059,10 +1056,7 @@ def get_empty_frame(data) -> DataFrame:
             )
             sparse_series.append(Series(data=sarr, index=index, name=col))
 
-        out = concat(sparse_series, axis=1, copy=False)
-        # TODO: overload concat with Literal for axis
-        out = cast(DataFrame, out)
-        return out
+        return concat(sparse_series, axis=1, copy=False)
 
     else:
         # take on axis=1 + transpose to ensure ndarray layout is column-major