From 286dfb126e4a295df8e2a2e31d765ef44170dd13 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lo=C3=AFc=20Diridollou?= <loic.diridollou@gmail.com>
Date: Fri, 13 Jun 2025 18:14:53 -0400
Subject: [PATCH 1/2] GH1248/GH1249 DataFrame.assign and
 MultiIndex.from_product Series/Index

---
 pandas-stubs/_typing.pyi            |  5 ++++-
 pandas-stubs/core/indexes/multi.pyi |  2 +-
 tests/test_frame.py                 | 15 ++++++++++++++-
 tests/test_indexes.py               |  4 ++++
 4 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi
index c404fc746..e231b669d 100644
--- a/pandas-stubs/_typing.pyi
+++ b/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
diff --git a/pandas-stubs/core/indexes/multi.pyi b/pandas-stubs/core/indexes/multi.pyi
index 506520208..f0b00ec39 100644
--- a/pandas-stubs/core/indexes/multi.pyi
+++ b/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: ...
diff --git a/tests/test_frame.py b/tests/test_frame.py
index eb0da9319..91b253987 100644
--- a/tests/test_frame.py
+++ b/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]})
diff --git a/tests/test_indexes.py b/tests/test_indexes.py
index 14bbc762e..e9be09aaf 100644
--- a/tests/test_indexes.py
+++ b/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"])

From c11881059f9cf27eacf99e4bca795156a5dab418 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lo=C3=AFc=20Diridollou?= <loic.diridollou@gmail.com>
Date: Fri, 13 Jun 2025 19:38:08 -0400
Subject: [PATCH 2/2] GH1248/GH1249 PR feedback

---
 pandas-stubs/_typing.pyi |  2 ++
 tests/test_frame.py      | 10 ++++++++--
 tests/test_indexes.py    | 14 ++++++++++++--
 3 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi
index e231b669d..7ee8e3f95 100644
--- a/pandas-stubs/_typing.pyi
+++ b/pandas-stubs/_typing.pyi
@@ -1016,6 +1016,8 @@ IntoColumn: TypeAlias = (
     AnyArrayLike
     | Scalar
     | Callable[[DataFrame], AnyArrayLike | Scalar | list[Scalar] | range]
+    | list[Scalar]
+    | range
     | None
 )
 
diff --git a/tests/test_frame.py b/tests/test_frame.py
index 91b253987..094a45462 100644
--- a/tests/test_frame.py
+++ b/tests/test_frame.py
@@ -45,7 +45,9 @@
 )
 import xarray as xr
 
-from pandas._typing import Scalar
+from pandas._typing import (
+    Scalar,
+)
 
 from tests import (
     PD_LTE_23,
@@ -315,11 +317,15 @@ def test_types_assign() -> None:
     df = pd.DataFrame({"a": [1, 2, 3]})
     check(
         assert_type(
-            df.assign(b=lambda df: range(len(df)), c=lambda df: [10, 20, 30]),
+            df.assign(b=lambda df: range(len(df)), c=lambda _: [10, 20, 30]),
             pd.DataFrame,
         ),
         pd.DataFrame,
     )
+    check(
+        assert_type(df.assign(b=range(len(df)), c=[10, 20, 30]), pd.DataFrame),
+        pd.DataFrame,
+    )
 
 
 def test_assign() -> None:
diff --git a/tests/test_indexes.py b/tests/test_indexes.py
index e9be09aaf..2b4480751 100644
--- a/tests/test_indexes.py
+++ b/tests/test_indexes.py
@@ -62,8 +62,18 @@ def test_index_astype() -> None:
     )
 
     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])])
+    check(
+        assert_type(
+            pd.MultiIndex.from_product([["x", "y"], df.columns]), pd.MultiIndex
+        ),
+        pd.MultiIndex,
+    )
+    check(
+        assert_type(
+            pd.MultiIndex.from_product([["x", "y"], pd.Series([1, 2])]), pd.MultiIndex
+        ),
+        pd.MultiIndex,
+    )
 
 
 def test_multiindex_get_level_values() -> None: