diff --git a/pandas-stubs/core/reshape/concat.pyi b/pandas-stubs/core/reshape/concat.pyi index 7cbcf65f2..90e4c6fa4 100644 --- a/pandas-stubs/core/reshape/concat.pyi +++ b/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], *, diff --git a/tests/test_pandas.py b/tests/test_pandas.py index eb4a10b95..c2f8f0773 100644 --- a/tests/test_pandas.py +++ b/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"}},