concat: accept Iterable as keys (#639)

Author: fostermaier

pandas-stubs/core/reshape/concat.pyi

@@ -29,7 +29,7 @@ def concat(
     axis: AxisIndex = ...,
     join: Literal["inner", "outer"] = ...,
     ignore_index: bool = ...,
-    keys: list[HashableT2] = ...,
+    keys: Iterable[HashableT2] = ...,
     levels: Sequence[list[HashableT3] | tuple[HashableT3, ...]] = ...,
     names: list[HashableT4] = ...,
     verify_integrity: bool = ...,
@@ -43,7 +43,7 @@ def concat(
     axis: AxisIndex = ...,
     join: Literal["inner", "outer"] = ...,
     ignore_index: bool = ...,
-    keys: list[HashableT2] = ...,
+    keys: Iterable[HashableT2] = ...,
     levels: Sequence[list[HashableT3] | tuple[HashableT3, ...]] = ...,
     names: list[HashableT4] = ...,
     verify_integrity: bool = ...,
@@ -57,7 +57,7 @@ def concat(
     axis: AxisColumn,
     join: Literal["inner", "outer"] = ...,
     ignore_index: bool = ...,
-    keys: list[HashableT2] = ...,
+    keys: Iterable[HashableT2] = ...,
     levels: Sequence[list[HashableT3] | tuple[HashableT3, ...]] = ...,
     names: list[HashableT4] = ...,
     verify_integrity: bool = ...,

tests/test_pandas.py

@@ -166,6 +166,16 @@ def test_concat_args() -> None:
         ),
         pd.DataFrame,
     )
+
+    df_dict = {"df1": df, "df2": df2}
+    check(
+        assert_type(
+            pd.concat(df_dict.values(), keys=df_dict.keys()),
+            pd.DataFrame,
+        ),
+        pd.DataFrame,
+    )
+
     check(
         assert_type(pd.concat([df, df2], ignore_index=True), pd.DataFrame), pd.DataFrame
     )