handle pyright inferring mixed float and int lists (#1097)

pandas-dev · Jan 14, 2025 · 0e22870 · 0e22870
1 parent 3547cda
commit 0e22870
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 1 deletion.
diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi
@@ -343,7 +343,7 @@ class Series(IndexOpsMixin[S1], NDFrame):
         copy: bool = ...,
     ) -> IntervalSeries[_OrderableT]: ...
     @overload
-    def __new__(
+    def __new__(  # type: ignore[overload-overlap]
         cls,
         data: Scalar | _ListLike | dict[HashableT1, Any] | None,
         index: Axes | None = ...,
@@ -353,6 +353,46 @@ class Series(IndexOpsMixin[S1], NDFrame):
         copy: bool = ...,
     ) -> Self: ...
     @overload
+    def __new__(  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
+        cls,
+        data: Sequence[bool],
+        index: Axes | None = ...,
+        *,
+        dtype: Dtype = ...,
+        name: Hashable = ...,
+        copy: bool = ...,
+    ) -> Series[bool]: ...
+    @overload
+    def __new__(  # type: ignore[overload-overlap]
+        cls,
+        data: Sequence[int],
+        index: Axes | None = ...,
+        *,
+        dtype: Dtype = ...,
+        name: Hashable = ...,
+        copy: bool = ...,
+    ) -> Series[int]: ...
+    @overload
+    def __new__(
+        cls,
+        data: Sequence[float],
+        index: Axes | None = ...,
+        *,
+        dtype: Dtype = ...,
+        name: Hashable = ...,
+        copy: bool = ...,
+    ) -> Series[float]: ...
+    @overload
+    def __new__(  # type: ignore[overload-cannot-match] # pyright: ignore[reportOverlappingOverload]
+        cls,
+        data: Sequence[int | float],
+        index: Axes | None = ...,
+        *,
+        dtype: Dtype = ...,
+        name: Hashable = ...,
+        copy: bool = ...,
+    ) -> Series[float]: ...
+    @overload
     def __new__(
         cls,
         data: S1 | _ListLike[S1] | dict[HashableT1, S1] | dict_keys[S1, Any],

diff --git a/tests/test_series.py b/tests/test_series.py
@@ -3539,3 +3539,16 @@ def test_series_dict() -> None:
         pd.Series,
         str,
     )
+
+
+def test_series_int_float() -> None:
+    # pyright infers mixtures of int and float in a list as list[int | float]
+    check(assert_type(pd.Series([1, 2, 3]), "pd.Series[int]"), pd.Series, np.integer)
+    check(
+        assert_type(pd.Series([1.0, 2.0, 3.0]), "pd.Series[float]"),
+        pd.Series,
+        np.float64,
+    )
+    check(
+        assert_type(pd.Series([1, 2.0, 3]), "pd.Series[float]"), pd.Series, np.float64
+    )