overload concat function

Author: topper-123

pandas/core/algorithms.py

@@ -3,7 +3,7 @@
 intended for public consumption
 """
 from textwrap import dedent
-from typing import Dict, Optional, TYPE_CHECKING, Tuple, Union
+from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union
 from warnings import catch_warnings, simplefilter, warn
 
 import numpy as np

pandas/core/reshape/concat.py

@@ -2,7 +2,7 @@
 concat routines
 """
 
-from typing import List, Union
+from typing import List, Mapping, Sequence, Union, overload
 
 import numpy as np
 
@@ -25,9 +25,28 @@
 # ---------------------------------------------------------------------
 # Concatenate DataFrame objects
 
+FrameOrSeriesUnion = Union["DataFrame", "Series"]
 
+
+@overload
+def concat(
+    objs: Union[Sequence["DataFrame"], Mapping[str, "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,
+) -> "DataFrame":
+    ...
+
+
+@overload
 def concat(
-    objs,
+    objs: Union[Sequence[FrameOrSeriesUnion], Mapping[str, FrameOrSeriesUnion]],
     axis=0,
     join: str = "outer",
     ignore_index: bool = False,
@@ -37,6 +56,21 @@ def concat(
     verify_integrity: bool = False,
     sort: bool = False,
     copy: bool = True,
+) -> Union["DataFrame", "Series"]:
+    ...
+
+
+def concat(
+    objs: Union[Sequence[FrameOrSeriesUnion], Mapping[str, FrameOrSeriesUnion]],
+    axis=0,
+    join="outer",
+    ignore_index: bool = False,
+    keys=None,
+    levels=None,
+    names=None,
+    verify_integrity: bool = False,
+    sort: bool = False,
+    copy: bool = True,
 ) -> Union["DataFrame", "Series"]:
     """
     Concatenate pandas objects along a particular axis with optional set logic

pandas/core/reshape/pivot.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Callable, Dict, Tuple, Union
+from typing import TYPE_CHECKING, Callable, Dict, List, Tuple, Union
 
 import numpy as np
 
@@ -40,7 +40,7 @@ def pivot_table(
     columns = _convert_by(columns)
 
     if isinstance(aggfunc, list):
-        pieces = []
+        pieces: List[DataFrame] = []
         keys = []
         for func in aggfunc:
             table = pivot_table(
@@ -58,9 +58,7 @@ def pivot_table(
             pieces.append(table)
             keys.append(getattr(func, "__name__", func))
 
-        result = concat(pieces, keys=keys, axis=1)
-        assert isinstance(result, DataFrame)
-        return result
+        return concat(pieces, keys=keys, axis=1)
 
     keys = index + columns
 

pandas/core/reshape/reshape.py

@@ -922,9 +922,7 @@ def check_len(item, name):
                 dtype=dtype,
             )
             with_dummies.append(dummy)
-        concatted = concat(with_dummies, axis=1)
-        assert isinstance(concatted, DataFrame)
-        result = concatted
+        result = concat(with_dummies, axis=1)
     else:
         result = _get_dummies_1d(
             data,

pandas/io/formats/format.py

@@ -260,6 +260,7 @@ def __init__(
         self._chk_truncate()
 
     def _chk_truncate(self) -> None:
+        from pandas.core.series import Series
         from pandas.core.reshape.concat import concat
 
         self.tr_row_num: Optional[int]
@@ -678,11 +679,9 @@ def _chk_truncate(self) -> None:
                 col_num = max_cols
             else:
                 col_num = max_cols_adj // 2
-                concatted = concat(
+                frame = concat(
                     (frame.iloc[:, :col_num], frame.iloc[:, -col_num:]), axis=1
                 )
-                assert isinstance(concatted, DataFrame)
-                frame = concatted
                 # truncate formatter
                 if isinstance(self.formatters, (list, tuple)):
                     truncate_fmt = self.formatters
@@ -699,9 +698,7 @@ def _chk_truncate(self) -> None:
                 frame = frame.iloc[:max_rows, :]
             else:
                 row_num = max_rows_adj // 2
-                concatted = concat((frame.iloc[:row_num, :], frame.iloc[-row_num:, :]))
-                assert isinstance(concatted, DataFrame)
-                frame = concatted
+                frame = concat((frame.iloc[:row_num, :], frame.iloc[-row_num:, :]))
             self.tr_row_num = row_num
         else:
             self.tr_row_num = None

pandas/io/pytables.py

@@ -4432,9 +4432,7 @@ def read(
         if len(frames) == 1:
             df = frames[0]
         else:
-            concatted = concat(frames, axis=1)
-            assert isinstance(concatted, DataFrame)
-            df = concatted
+            df = concat(frames, axis=1)
 
         selection = Selection(self, where=where, start=start, stop=stop)
         # apply the selection filters & axis orderings