Add API reference documentation & Implement insert_level using levels/codes operations

Author: cloudboat

doc/source/reference/indexing.rst

@@ -294,6 +294,7 @@ MultiIndex components
    MultiIndex.copy
    MultiIndex.append
    MultiIndex.truncate
+   MultiIndex.insert_level
 
 MultiIndex selecting
 ~~~~~~~~~~~~~~~~~~~~

pandas/core/indexes/multi.py

@@ -2710,19 +2710,21 @@ def reorder_levels(self, order) -> MultiIndex:
         result = self._reorder_ilevels(order)
         return result
 
-    def insert_level(self, position: int, value, name=None):
+    def insert_level(
+        self, position: int, value, name: Hashable = lib.no_default
+    ) -> MultiIndex:
         """
         Insert a new level at the specified position in the MultiIndex.
 
         Parameters
         ----------
         position : int
             The position at which to insert the new level (0-based).
-        value : scalar or array-like
-            Value(s) to use for the new level. If scalar, broadcast to all items.
-            If array-like, length must match the length of the index.
-        name : object, optional
-            Name for the new level.
+            Must be between 0 and nlevels (inclusive).
+        value : array-like
+            Values to use for the new level. Length must match the length of the index.
+        name : Hashable, default lib.no_default
+            Name for the new level. If not provided, the new level will have no name.
 
         Returns
         -------
@@ -2732,37 +2734,40 @@ def insert_level(self, position: int, value, name=None):
         Examples
         --------
         >>> idx = pd.MultiIndex.from_tuples([("A", 1), ("B", 2)])
-        >>> idx.insert_level(0, "new_value")
+        >>> idx.insert_level(0, ["new_value", "new_value"])
         MultiIndex([('new_value', 'A', 1), ('new_value', 'B', 2)], ...)
 
         >>> idx.insert_level(1, ["X", "Y"])
         MultiIndex([('A', 'X', 1), ('B', 'Y', 2)], ...)
-
-        >>> idx.insert_level(0, "new_val", name="new_level")
-        MultiIndex([('new_val', 'A', 1), ('new_val', 'B', 2)], ...)
         """
         if not isinstance(position, int):
             raise TypeError("position must be an integer")
 
         if position < 0 or position > self.nlevels:
             raise ValueError(f"position must be between 0 and {self.nlevels}")
 
+        if name is lib.no_default:
+            name = None
+
         if not hasattr(value, "__iter__") or isinstance(value, str):
-            value = [value] * len(self)
-        else:
-            value = list(value)
-            if len(value) != len(self):
-                raise ValueError("Length of values must match length of index")
+            raise TypeError("value must be an array-like object")
 
-        new_tuples = []
+        value = list(value)
+        if len(value) != len(self):
+            raise ValueError("Length of values must match length of index")
 
-        for i, tup in enumerate(self):
-            new_tuple = tup[:position] + (value[i],) + tup[position:]
-            new_tuples.append(new_tuple)
+        new_level = Index(value)
+        new_codes_for_level = new_level.get_indexer(value)
 
+        new_levels = self.levels[:position] + [new_level] + self.levels[position:]
+        new_codes = (
+            self.codes[:position] + [new_codes_for_level] + self.codes[position:]
+        )
         new_names = self.names[:position] + [name] + self.names[position:]
 
-        return MultiIndex.from_tuples(new_tuples, names=new_names)
+        return MultiIndex(
+            levels=new_levels, codes=new_codes, names=new_names, verify_integrity=False
+        )
 
     def _reorder_ilevels(self, order) -> MultiIndex:
         if len(order) != self.nlevels:

pandas/tests/indexes/multi/test_constructors.py

@@ -870,14 +870,3 @@ def test_dtype_representation(using_infer_string):
         dtype=object,
     )
     tm.assert_series_equal(result, expected)
-
-
-def test_insert_level_integration():
-    idx = MultiIndex.from_tuples([("A", 1), ("B", 2)])
-
-    df = pd.DataFrame({"data": [10, 20]}, index=idx)
-    new_idx = idx.insert_level(0, "group1")
-    df_new = df.set_index(new_idx)
-
-    assert df_new.index.nlevels == 3
-    assert len(df_new) == 2

pandas/tests/indexes/multi/test_insert_level.py

@@ -4,146 +4,148 @@
 import pandas._testing as tm
 
 
-class TestMultiIndexInsertLevel:
-    @pytest.mark.parametrize(
-        "position, value, name, expected_tuples, expected_names",
-        [
-            (
-                0,
-                "new_value",
-                None,
-                [("new_value", "A", 1), ("new_value", "B", 2), ("new_value", "C", 3)],
-                [None, "level1", "level2"],
-            ),
-            (
-                1,
-                "middle",
-                None,
-                [("A", "middle", 1), ("B", "middle", 2), ("C", "middle", 3)],
-                ["level1", None, "level2"],
-            ),
-            (
-                0,
-                "new_val",
-                "new_level",
-                [("new_val", "A", 1), ("new_val", "B", 2), ("new_val", "C", 3)],
-                ["new_level", "level1", "level2"],
-            ),
-            (
-                1,
-                "middle",
-                "custom_name",
-                [("A", "middle", 1), ("B", "middle", 2), ("C", "middle", 3)],
-                ["level1", "custom_name", "level2"],
-            ),
-            (
-                0,
-                "start",
-                None,
-                [("start", "A", 1), ("start", "B", 2), ("start", "C", 3)],
-                [None, "level1", "level2"],
-            ),
-            (
-                2,
-                "end",
-                None,
-                [("A", 1, "end"), ("B", 2, "end"), ("C", 3, "end")],
-                ["level1", "level2", None],
-            ),
-            (
-                1,
-                100,
-                None,
-                [("A", 100, 1), ("B", 100, 2), ("C", 100, 3)],
-                ["level1", None, "level2"],
-            ),
-            (
-                1,
-                1.5,
-                None,
-                [("A", 1.5, 1), ("B", 1.5, 2), ("C", 1.5, 3)],
-                ["level1", None, "level2"],
-            ),
-            (
-                1,
-                None,
-                None,
-                [("A", None, 1), ("B", None, 2), ("C", None, 3)],
-                ["level1", None, "level2"],
-            ),
-            (
-                1,
-                ["X", "Y", "Z"],
-                None,
-                [("A", "X", 1), ("B", "Y", 2), ("C", "Z", 3)],
-                ["level1", None, "level2"],
-            ),
-            (
-                0,
-                "",
-                "empty_string",
-                [("", "A", 1), ("", "B", 2), ("", "C", 3)],
-                ["empty_string", "level1", "level2"],
-            ),
-            (
-                1,
-                True,
-                None,
-                [("A", True, 1), ("B", True, 2), ("C", True, 3)],
-                ["level1", None, "level2"],
-            ),
-        ],
+@pytest.mark.parametrize(
+    "position, value, name, expected_tuples, expected_names",
+    [
+        (
+            0,
+            ["new_value"] * 3,
+            None,
+            [("new_value", "A", 1), ("new_value", "B", 2), ("new_value", "C", 3)],
+            [None, "level1", "level2"],
+        ),
+        (
+            1,
+            ["middle"] * 3,
+            None,
+            [("A", "middle", 1), ("B", "middle", 2), ("C", "middle", 3)],
+            ["level1", None, "level2"],
+        ),
+        (
+            0,
+            ["new_val"] * 3,
+            "new_level",
+            [("new_val", "A", 1), ("new_val", "B", 2), ("new_val", "C", 3)],
+            ["new_level", "level1", "level2"],
+        ),
+        (
+            1,
+            ["middle"] * 3,
+            "custom_name",
+            [("A", "middle", 1), ("B", "middle", 2), ("C", "middle", 3)],
+            ["level1", "custom_name", "level2"],
+        ),
+        (
+            0,
+            ["start"] * 3,
+            None,
+            [("start", "A", 1), ("start", "B", 2), ("start", "C", 3)],
+            [None, "level1", "level2"],
+        ),
+        (
+            2,
+            ["end"] * 3,
+            None,
+            [("A", 1, "end"), ("B", 2, "end"), ("C", 3, "end")],
+            ["level1", "level2", None],
+        ),
+        (
+            1,
+            [100, 100, 100],
+            None,
+            [("A", 100, 1), ("B", 100, 2), ("C", 100, 3)],
+            ["level1", None, "level2"],
+        ),
+        (
+            1,
+            [1.5, 1.5, 1.5],
+            None,
+            [("A", 1.5, 1), ("B", 1.5, 2), ("C", 1.5, 3)],
+            ["level1", None, "level2"],
+        ),
+        (
+            1,
+            [None, None, None],
+            None,
+            [("A", None, 1), ("B", None, 2), ("C", None, 3)],
+            ["level1", None, "level2"],
+        ),
+        (
+            1,
+            ["X", "Y", "Z"],
+            None,
+            [("A", "X", 1), ("B", "Y", 2), ("C", "Z", 3)],
+            ["level1", None, "level2"],
+        ),
+        (
+            0,
+            [""] * 3,
+            "empty_string",
+            [("", "A", 1), ("", "B", 2), ("", "C", 3)],
+            ["empty_string", "level1", "level2"],
+        ),
+        (
+            1,
+            [True, True, True],
+            None,
+            [("A", True, 1), ("B", True, 2), ("C", True, 3)],
+            ["level1", None, "level2"],
+        ),
+    ],
+)
+def test_insert_level_basic(position, value, name, expected_tuples, expected_names):
+    simple_idx = pd.MultiIndex.from_tuples(
+        [("A", 1), ("B", 2), ("C", 3)], names=["level1", "level2"]
     )
-    def test_insert_level_basic(
-        self, position, value, name, expected_tuples, expected_names
-    ):
-        simple_idx = pd.MultiIndex.from_tuples(
-            [("A", 1), ("B", 2), ("C", 3)], names=["level1", "level2"]
-        )
-
-        result = simple_idx.insert_level(position, value, name=name)
-        expected = pd.MultiIndex.from_tuples(expected_tuples, names=expected_names)
-        tm.assert_index_equal(result, expected)
-
-    @pytest.mark.parametrize(
-        "position, value, expected_error",
-        [
-            (5, "invalid", "position must be between"),
-            (-1, "invalid", "position must be between"),
-            (1, ["too", "few"], "Length of values must match"),
-            (3, "value", "position must be between"),
-        ],
+
+    result = simple_idx.insert_level(position, value, name=name)
+    expected = pd.MultiIndex.from_tuples(expected_tuples, names=expected_names)
+    tm.assert_index_equal(result, expected)
+
+
+@pytest.mark.parametrize(
+    "position, value, expected_error",
+    [
+        (5, ["invalid"] * 3, "position must be between"),
+        (-1, ["invalid"] * 3, "position must be between"),
+        (1, ["too", "few"], "Length of values must match"),
+        (3, ["value"] * 3, "position must be between"),
+        (0, "scalar_value", "value must be an array-like object"),
+    ],
+)
+def test_insert_level_error_cases(position, value, expected_error):
+    simple_idx = pd.MultiIndex.from_tuples(
+        [("A", 1), ("B", 2), ("C", 3)], names=["level1", "level2"]
+    )
+
+    with pytest.raises(ValueError, match=expected_error):
+        simple_idx.insert_level(position, value)
+
+
+def test_insert_level_preserves_original():
+    simple_idx = pd.MultiIndex.from_tuples(
+        [("A", 1), ("B", 2), ("C", 3)], names=["level1", "level2"]
     )
-    def test_insert_level_error_cases(self, position, value, expected_error):
-        simple_idx = pd.MultiIndex.from_tuples(
-            [("A", 1), ("B", 2), ("C", 3)], names=["level1", "level2"]
-        )
 
-        with pytest.raises(ValueError, match=expected_error):
-            simple_idx.insert_level(position, value)
+    original = simple_idx.copy()
+    simple_idx.insert_level(1, ["temp"] * 3)
 
-    def test_insert_level_preserves_original(self):
-        simple_idx = pd.MultiIndex.from_tuples(
-            [("A", 1), ("B", 2), ("C", 3)], names=["level1", "level2"]
-        )
+    tm.assert_index_equal(original, simple_idx)
 
-        original = simple_idx.copy()
-        simple_idx.insert_level(1, "temp")
 
-        tm.assert_index_equal(original, simple_idx)
+def test_insert_level_empty_index():
+    empty_idx = pd.MultiIndex.from_tuples([], names=["level1", "level2"])
 
-    def test_insert_level_empty_index(self):
-        empty_idx = pd.MultiIndex.from_tuples([], names=["level1", "level2"])
+    result = empty_idx.insert_level(0, [])
+    expected = pd.MultiIndex.from_tuples([], names=[None, "level1", "level2"])
+    tm.assert_index_equal(result, expected)
 
-        result = empty_idx.insert_level(0, [])
-        expected = pd.MultiIndex.from_tuples([], names=[None, "level1", "level2"])
-        tm.assert_index_equal(result, expected)
 
-    def test_insert_level_single_element(self):
-        single_idx = pd.MultiIndex.from_tuples([("A", 1)], names=["level1", "level2"])
+def test_insert_level_single_element():
+    single_idx = pd.MultiIndex.from_tuples([("A", 1)], names=["level1", "level2"])
 
-        result = single_idx.insert_level(1, "middle")
-        expected = pd.MultiIndex.from_tuples(
-            [("A", "middle", 1)], names=["level1", None, "level2"]
-        )
-        tm.assert_index_equal(result, expected)
+    result = single_idx.insert_level(1, ["middle"])
+    expected = pd.MultiIndex.from_tuples(
+        [("A", "middle", 1)], names=["level1", None, "level2"]
+    )
+    tm.assert_index_equal(result, expected)