merge

Author: yangdanny97

pandas-stubs/core/reshape/concat.pyi

@@ -15,6 +15,7 @@ from pandas import (
 from typing_extensions import Never
 
 from pandas._typing import (
+    S2,
     Axis,
     AxisIndex,
     HashableT1,
@@ -38,7 +39,21 @@ def concat(  # type: ignore[overload-overlap]
     copy: bool = ...,
 ) -> DataFrame: ...
 @overload
-def concat(  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
+def concat(  # pyright: ignore[reportOverlappingOverload]
+    objs: Iterable[Series[S2]],
+    *,
+    axis: AxisIndex = ...,
+    join: Literal["inner", "outer"] = ...,
+    ignore_index: bool = ...,
+    keys: Iterable[HashableT2] = ...,
+    levels: Sequence[list[HashableT3] | tuple[HashableT3, ...]] = ...,
+    names: list[HashableT4] | None = ...,
+    verify_integrity: bool = ...,
+    sort: bool = ...,
+    copy: bool = ...,
+) -> Series[S2]: ...
+@overload
+def concat(  # type: ignore[overload-overlap]
     objs: Iterable[Series] | Mapping[HashableT1, Series],
     *,
     axis: AxisIndex = ...,

tests/series/test_series.py

@@ -63,6 +63,10 @@
 from pandas.io.formats.format import EngFormatter
 from pandas.tseries.offsets import (
     BaseOffset,
+    BDay,
+    BQuarterEnd,
+    MonthEnd,
+    Week,
     YearEnd,
 )
 
@@ -448,6 +452,7 @@ def test_types_sort_values_with_key() -> None:
 
 
 def test_types_shift() -> None:
+    """Test shift operator on series with different arguments."""
     s = pd.Series([1, 2, 3], index=pd.date_range("2020", periods=3))
     check(assert_type(s.shift(), pd.Series), pd.Series, np.floating)
     check(
@@ -457,6 +462,11 @@ def test_types_shift() -> None:
     )
     check(assert_type(s.shift(-1, fill_value=0), pd.Series), pd.Series, np.integer)
     check(assert_type(s.shift(freq="1D"), pd.Series), pd.Series, np.integer)
+    check(assert_type(s.shift(freq=BDay(1)), pd.Series), pd.Series, np.integer)
+    check(assert_type(s.shift(freq=BQuarterEnd(5)), pd.Series), pd.Series, np.integer)
+    check(assert_type(s.shift(freq=MonthEnd(3)), pd.Series), pd.Series, np.integer)
+    check(assert_type(s.shift(freq=Week(4)), pd.Series), pd.Series, np.integer)
+    check(assert_type(s.shift(freq=YearEnd(2)), pd.Series), pd.Series, np.integer)
 
 
 def test_series_pct_change() -> None:

tests/test_frame.py

@@ -60,6 +60,13 @@
 from pandas.io.formats.format import EngFormatter
 from pandas.io.formats.style import Styler
 from pandas.io.parsers import TextFileReader
+from pandas.tseries.offsets import (
+    BDay,
+    BQuarterEnd,
+    MonthEnd,
+    Week,
+    YearEnd,
+)
 
 if TYPE_CHECKING:
     from pandas.core.frame import _PandasNamedTuple
@@ -603,6 +610,11 @@ def test_types_shift() -> None:
     check(assert_type(df.shift(1), pd.DataFrame), pd.DataFrame)
     check(assert_type(df.shift(-1), pd.DataFrame), pd.DataFrame)
     check(assert_type(df.shift(freq="1D"), pd.DataFrame), pd.DataFrame)
+    check(assert_type(df.shift(freq=BDay(1)), pd.DataFrame), pd.DataFrame)
+    check(assert_type(df.shift(freq=BQuarterEnd(5)), pd.DataFrame), pd.DataFrame)
+    check(assert_type(df.shift(freq=MonthEnd(3)), pd.DataFrame), pd.DataFrame)
+    check(assert_type(df.shift(freq=Week(4)), pd.DataFrame), pd.DataFrame)
+    check(assert_type(df.shift(freq=YearEnd(2)), pd.DataFrame), pd.DataFrame)
 
 
 def test_types_rank() -> None:
@@ -2523,8 +2535,8 @@ def test_types_regressions() -> None:
     ts1 = pd.concat([s1, s2], axis=0)
     ts2 = pd.concat([s1, s2])
 
-    check(assert_type(ts1, pd.Series), pd.Series)
-    check(assert_type(ts2, pd.Series), pd.Series)
+    check(assert_type(ts1, "pd.Series[int]"), pd.Series, np.integer)
+    check(assert_type(ts2, "pd.Series[int]"), pd.Series, np.integer)
 
     # https://github.com/microsoft/python-type-stubs/issues/110
     check(assert_type(pd.Timestamp("2021-01-01"), pd.Timestamp), datetime.date)

tests/test_pandas.py

@@ -132,25 +132,30 @@ def test_types_concat() -> None:
     s = pd.Series([0, 1, -10])
     s2 = pd.Series([7, -5, 10])
 
-    check(assert_type(pd.concat([s, s2]), pd.Series), pd.Series)
+    check(assert_type(pd.concat([s, s2]), "pd.Series[int]"), pd.Series, np.integer)
     check(assert_type(pd.concat([s, s2], axis=1), pd.DataFrame), pd.DataFrame)
     check(
-        assert_type(pd.concat([s, s2], keys=["first", "second"], sort=True), pd.Series),
+        assert_type(
+            pd.concat([s, s2], keys=["first", "second"], sort=True), "pd.Series[int]"
+        ),
         pd.Series,
+        np.integer,
     )
     check(
         assert_type(
             pd.concat([s, s2], keys=["first", "second"], names=["source", "row"]),
-            pd.Series,
+            "pd.Series[int]",
         ),
         pd.Series,
+        np.integer,
     )
     check(
         assert_type(
             pd.concat([s, s2], keys=["first", "second"], names=None),
-            pd.Series,
+            "pd.Series[int]",
         ),
         pd.Series,
+        np.integer,
     )
 
     # Depends on the axis
@@ -236,6 +241,37 @@ def test_types_concat() -> None:
     check(assert_type(pd.concat([pd.DataFrame(), data]), pd.DataFrame), pd.DataFrame)
 
 
+def test_concat_series_mixed_numeric() -> None:
+    """Test concatenation of Series with mixed numeric types.
+
+    Derived from test_types_concat."""
+    s = pd.Series([0, 1, -10])
+    s2 = pd.Series([7.0, -5, 10])
+
+    check(assert_type(pd.concat([s, s2]), pd.Series), pd.Series, np.floating)
+    check(
+        assert_type(pd.concat([s, s2], keys=["first", "second"], sort=True), pd.Series),
+        pd.Series,
+        np.floating,
+    )
+    check(
+        assert_type(
+            pd.concat([s, s2], keys=["first", "second"], names=["source", "row"]),
+            pd.Series,
+        ),
+        pd.Series,
+        np.floating,
+    )
+    check(
+        assert_type(
+            pd.concat([s, s2], keys=["first", "second"], names=None),
+            pd.Series,
+        ),
+        pd.Series,
+        np.floating,
+    )
+
+
 def test_concat_args() -> None:
     df = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]})
     df2 = pd.DataFrame(data={"col1": [10, 20], "col2": [30, 40]}, index=[2, 3])