overload concat function

Author: topper-123

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,12 @@
 # ---------------------------------------------------------------------
 # Concatenate DataFrame objects
 
+FrameOrSeriesUnion = Union["DataFrame", "Series"]
 
+
+@overload
 def concat(
-    objs,
+    objs: Union[Sequence["DataFrame"], Mapping[str, "DataFrame"]],
     axis=0,
     join: str = "outer",
     ignore_index: bool = False,
@@ -37,6 +40,33 @@ def concat(
     verify_integrity: bool = False,
     sort: bool = False,
     copy: bool = True,
+) -> "DataFrame":
+    ...
+@overload  # noqa: E302
+def concat(
+    objs: Union[Sequence[FrameOrSeriesUnion], Mapping[str, FrameOrSeriesUnion]],
+    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,
+) -> Union["DataFrame", "Series"]:
+    ...
+def concat(  # noqa: E302
+        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

@@ -678,11 +678,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 +697,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