GH1248/GH1249 DataFrame.assign and MultiIndex.from_product Series/Index

Author: loicdiridollou

pandas-stubs/_typing.pyi

@@ -1013,7 +1013,10 @@ TimeZones: TypeAlias = str | tzinfo | None | int
 
 # Evaluates to a DataFrame column in DataFrame.assign context.
 IntoColumn: TypeAlias = (
-    AnyArrayLike | Scalar | Callable[[DataFrame], AnyArrayLike | Scalar] | None
+    AnyArrayLike
+    | Scalar
+    | Callable[[DataFrame], AnyArrayLike | Scalar | list[Scalar] | range]
+    | None
 )
 
 DatetimeLike: TypeAlias = datetime.datetime | np.datetime64 | Timestamp

pandas-stubs/core/indexes/multi.pyi

@@ -55,7 +55,7 @@ class MultiIndex(Index[Any]):
     @classmethod
     def from_product(
         cls,
-        iterables: Sequence[SequenceNotStr[Hashable]],
+        iterables: Sequence[SequenceNotStr[Hashable] | pd.Series | pd.Index],
         sortorder: int | None = ...,
         names: SequenceNotStr[Hashable] = ...,
     ) -> Self: ...

tests/test_frame.py

@@ -305,9 +305,22 @@ def test_types_head_tail() -> None:
 
 def test_types_assign() -> None:
     df = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]})
-    df.assign(col3=lambda frame: frame.sum(axis=1))
+
+    check(
+        assert_type(df.assign(col3=lambda frame: frame.sum(axis=1)), pd.DataFrame),
+        pd.DataFrame,
+    )
     df["col3"] = df.sum(axis=1)
 
+    df = pd.DataFrame({"a": [1, 2, 3]})
+    check(
+        assert_type(
+            df.assign(b=lambda df: range(len(df)), c=lambda df: [10, 20, 30]),
+            pd.DataFrame,
+        ),
+        pd.DataFrame,
+    )
+
 
 def test_assign() -> None:
     df = pd.DataFrame({"a": [1, 2, 3], 1: [4, 5, 6]})

tests/test_indexes.py

@@ -61,6 +61,10 @@ def test_index_astype() -> None:
         pd.DataFrame,
     )
 
+    df = pd.DataFrame({"a": [1, 2, 3]})
+    mi = pd.MultiIndex.from_product([["x", "y"], df.columns])
+    mi = pd.MultiIndex.from_product([["x", "y"], pd.Series([1, 2])])
+
 
 def test_multiindex_get_level_values() -> None:
     mi = pd.MultiIndex.from_product([["a", "b"], ["c", "d"]], names=["ab", "cd"])