Fix: proper return types for MultiIndex.swaplevel & MultiIndex.union

pandas-stubs/_typing.pyi

@@ -904,6 +904,7 @@ SeriesDType: TypeAlias = (
     | datetime.datetime  # includes pd.Timestamp
     | datetime.timedelta  # includes pd.Timedelta
 )
+
 S1 = TypeVar("S1", bound=SeriesDType, default=Any)
 # Like S1, but without `default=Any`.
 S2 = TypeVar("S2", bound=SeriesDType)

pandas-stubs/core/indexes/multi.pyi

@@ -135,7 +135,7 @@ class MultiIndex(Index):
     def append(self, other): ...
     def repeat(self, repeats, axis=...): ...
     def drop(self, codes, level: Level | None = None, errors: str = "raise") -> Self: ...  # type: ignore[override] # pyright: ignore[reportIncompatibleMethodOverride]
-    def swaplevel(self, i: int = -2, j: int = -1): ...
+    def swaplevel(self, i: int = -2, j: int = -1) -> Self: ...
     def reorder_levels(self, order): ...
     def sortlevel(
         self,
@@ -163,3 +163,7 @@ class MultiIndex(Index):
     def insert(self, loc, item): ...
     def delete(self, loc): ...
     def isin(self, values, level=...) -> np_1darray[np.bool]: ...
+    @overload
+    def union(self, other: Self | list[tuple[Hashable, ...]], sort: bool | None = ...,) -> Self: ...
+    @overload
+    def union(self, other: Index | list[HashableT], sort: bool | None = ...,) -> Index: ...

tests/indexes/test_indexes.py

@@ -1613,3 +1613,31 @@ def test_to_series() -> None:
         np.complexfloating,
     )
     check(assert_type(Index(["1"]).to_series(), "pd.Series[str]"), pd.Series, str)
+
+
+def test_swaplevel_rettype() -> None:
+    """Test that swaplevel returns Self"""
+    mi = pd.MultiIndex.from_product([["a", "b"], [1, 2]], names=["let", "num"])
+    check(
+        assert_type(mi.swaplevel(0, 1), "pd.MultiIndex"),
+        pd.MultiIndex,
+    )
+
+
+def test_multiindex_union() -> None:
+    """Test that union returns MultiIndex on MultiIndex input"""
+    mi = pd.MultiIndex.from_product([["a", "b"], [1, 2]], names=["let", "num"])
+    mi2 = pd.MultiIndex.from_product([["a", "b"], [3, 4]], names=["let", "num"])
+
+    check(
+        assert_type(mi.union(mi2), "pd.MultiIndex"),
+        pd.MultiIndex,
+    )
+    check(
+        assert_type(mi.union(pd.Index([("c", 3), ("d", 4)])), "pd.Index"),
+        pd.Index,
+    )
+    check(
+        assert_type(mi.union([("c", 3), ("d", 4)]), "pd.MultiIndex"),
+        pd.MultiIndex,
+    )