Fix: proper return types for MultiIndex.swaplevel & MultiIndex.union (#1437)

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

* Index of and List of tuples in MultiIndex overriding

* MultiIndex union tests

* MultiIndex-Index union overload simplified

* MultiIndex test amended and swaplevel test added

* Simplified MultiIndex-Index union overload

* Reverted S1 changes

* amended MultiIndex union test and added swaplevel test

* Removed runtime-failing test

* Pre-commit refactoring

* Removed a MultiIndex-Index test

* Removed MultiIndex union overloading

* Separated swaplevel and union tests

Co-authored-by: Yi-Fan Wang <[email protected]>

* test for MultiIndex-Index union

* MultiIndex.union overload replaced with change of base return type

* removed MultiIndex.union overload

* Removed extra newline

Co-authored-by: Yi-Fan Wang <[email protected]>

* Added py.typed

* existing tests fixed

* tests fixed for runtime errors

---------

Co-authored-by: Yi-Fan Wang <[email protected]>

Author: realzachbrownjohn

pandas-stubs/core/indexes/base.pyi

@@ -405,7 +405,7 @@ class Index(IndexOpsMixin[S1], ElementOpsMixin[S1]):
     __bool__ = ...
     def union(
         self, other: list[HashableT] | Self, sort: bool | None = None
-    ) -> Index: ...
+    ) -> Self: ...
     def intersection(
         self, other: list[S1] | Self, sort: bool | None = False
     ) -> Self: ...

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,

pandas-stubs/core/indexes/range.pyi

@@ -80,7 +80,7 @@ class RangeIndex(_IndexSubclassBase[int, np.int64]):
     def all(self, *args: Any, **kwargs: Any) -> bool: ...
     def any(self, *args: Any, **kwargs: Any) -> bool: ...
     @final
-    def union(
+    def union(  # type: ignore[override]
         self, other: list[HashableT] | Index, sort: bool | None = None
     ) -> Index | Index[int] | RangeIndex: ...
     @overload  # type: ignore[override]

tests/indexes/test_indexes.py

@@ -302,12 +302,18 @@ def test_range_index_union() -> None:
 def test_index_union_sort() -> None:
     """Test sort argument in pd.Index.union GH1264."""
     check(
-        assert_type(pd.Index(["e", "f"]).union(["a", "b", "c"], sort=True), pd.Index),
+        assert_type(
+            pd.Index(["e", "f"]).union(["a", "b", "c"], sort=True), "pd.Index[str]"
+        ),
         pd.Index,
+        str,
     )
     check(
-        assert_type(pd.Index(["e", "f"]).union(["a", "b", "c"], sort=False), pd.Index),
+        assert_type(
+            pd.Index(["e", "f"]).union(["a", "b", "c"], sort=False), "pd.Index[str]"
+        ),
         pd.Index,
+        str,
     )
 
 
@@ -1520,3 +1526,18 @@ def test_to_series() -> None:
         np.complexfloating,
     )
     check(assert_type(Index(["1"]).to_series(), "pd.Series[str]"), pd.Series, str)
+
+
+def test_multiindex_union() -> None:
+    """Test that MultiIndex.union returns MultiIndex"""
+    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([("c", 3), ("d", 4)]), "pd.MultiIndex"), pd.MultiIndex)
+
+
+def test_multiindex_swaplevel() -> None:
+    """Test that MultiIndex.swaplevel returns MultiIndex"""
+    mi = pd.MultiIndex.from_product([["a", "b"], [1, 2]], names=["let", "num"])
+    check(assert_type(mi.swaplevel(0, 1), "pd.MultiIndex"), pd.MultiIndex)