GH1396 Add concat compatibility for subclasses of DataFrame (#1713)

* GH1396 Add concat compatibility for subclasses of DataFrame

* GH1396 Add concat compatibility for subclasses of DataFrame

* GH1396 Add concat compatibility for subclasses of DataFrame

Author: loicdiridollou

pandas-stubs/core/reshape/concat.pyi

@@ -6,6 +6,7 @@ from collections.abc import (
 from typing import (
     Literal,
     Never,
+    TypeVar,
     overload,
 )
 
@@ -23,6 +24,8 @@ from pandas._typing import (
     HashableT4,
 )
 
+DataFrameT0 = TypeVar("DataFrameT0", bound=DataFrame)
+
 @overload
 def concat(
     objs: Iterable[None] | Mapping[HashableT1, None],
@@ -63,6 +66,19 @@ def concat(  # type: ignore[overload-overlap]
     sort: bool = False,
 ) -> Series: ...
 @overload
+def concat(
+    objs: Iterable[DataFrameT0 | None],
+    *,
+    axis: Axis = 0,
+    join: Literal["inner", "outer"] = "outer",
+    ignore_index: bool = False,
+    keys: Iterable[HashableT2] | None = None,
+    levels: Sequence[list[HashableT3] | tuple[HashableT3, ...]] | None = None,
+    names: list[HashableT4] | None = None,
+    verify_integrity: bool = False,
+    sort: bool = False,
+) -> DataFrameT0: ...
+@overload
 def concat(
     objs: Iterable[NDFrame | None] | Mapping[HashableT1, NDFrame | None],
     *,

tests/test_pandas.py

@@ -404,6 +404,21 @@ def test_concat_args() -> None:
     )
 
 
+def test_frame_subclass_concat() -> None:
+    """Test concatenate subclass of DataFrame GH1396."""
+
+    class ChildDataFrame(pd.DataFrame):
+        @property
+        def _constructor(self) -> type[ChildDataFrame]:
+            return ChildDataFrame
+
+    cdf1 = ChildDataFrame(data={"a": [0]})
+    cdf2 = ChildDataFrame(data={"a": [1]})
+
+    cdf = pd.concat([cdf1, cdf2], ignore_index=True)
+    check(assert_type(cdf, ChildDataFrame), ChildDataFrame)
+
+
 def test_types_json_normalize() -> None:
     data1: list[dict[str, Any]] = [
         {"id": 1, "name": {"first": "Coleen", "last": "Volk"}},