enh(assignlevel): Add ignore_index argument (#41)

CHANGELOG.rst

@@ -3,6 +3,12 @@
 Changelog
 =========
 
+v0.3.1 (2023-09-18)
+------------------------------------------------------------
+* The new :func:`~core.assignlevel` argument `ignore_index=True` prevents the
+  dataframe and series alignment which became the default in v0.3 (yesterday),
+  since there are valid use cases of the old behaviour
+
 v0.3 (2023-09-17)
 ------------------------------------------------------------
 * **BREAKING** :func:`~core.assignlevel` aligns :class:`~pandas.Series` and

src/pandas_indexing/accessors.py

@@ -55,9 +55,17 @@ def assign(
         frame: Optional[Data] = None,
         order: bool = False,
         axis: Axis = 0,
+        ignore_index: bool = False,
         **labels: Any,
     ) -> Union[DataFrame, Series, MultiIndex]:
-        return assignlevel(self._obj, frame=frame, order=order, axis=axis, **labels)
+        return assignlevel(
+            self._obj,
+            frame=frame,
+            order=order,
+            axis=axis,
+            ignore_index=ignore_index,
+            **labels,
+        )
 
     @doc(extractlevel, index_or_data="")
     def extract(

src/pandas_indexing/core.py

@@ -31,7 +31,11 @@
 
 
 def _assignlevel(
-    index: Index, frame: Optional[Data] = None, order: bool = False, **labels: Any
+    index: Index,
+    frame: Optional[Data] = None,
+    order: bool = False,
+    ignore_index: bool = False,
+    **labels: Any,
 ) -> MultiIndex:
     if isinstance(index, MultiIndex):
         new_levels = list(index.levels)
@@ -43,11 +47,18 @@ def _assignlevel(
         new_codes = [level.get_indexer(index)]
         new_names = [index.name]
 
-    def maybe_align_with_index(df_or_ser):
-        if isinstance(df_or_ser, (DataFrame, Series)):
-            return semijoin(df_or_ser, index, how="right", fail_on_reorder=True)
+    if ignore_index:
+
+        def maybe_align_with_index(df_or_ser):
+            return df_or_ser
 
-        return df_or_ser
+    else:
+
+        def maybe_align_with_index(df_or_ser):
+            if isinstance(df_or_ser, (DataFrame, Series)):
+                return semijoin(df_or_ser, index, how="right", fail_on_reorder=True)
+
+            return df_or_ser
 
     labels = {level: maybe_align_with_index(lbls) for level, lbls in labels.items()}
 
@@ -97,6 +108,7 @@ def assignlevel(
     frame: Optional[Data] = None,
     order: bool = False,
     axis: Axis = 0,
+    ignore_index: bool = False,
     **labels: Any,
 ) -> T:
     """Add or overwrite levels on a multiindex.
@@ -110,6 +122,8 @@ def assignlevel(
         Level names in desired order or False, by default False
     axis : {{0, 1, "index", "columns"}}, default 0
         Axis where to update multiindex
+    ignore_index : bool, optional
+        If true, dataframes or series are not index aligned
     **labels
         Labels for each new index level
 
@@ -119,10 +133,14 @@ def assignlevel(
         Series or DataFrame with changed index or new MultiIndex
     """
     if isinstance(df, Index):
-        return _assignlevel(df, frame=frame, order=order, **labels)
+        return _assignlevel(
+            df, frame=frame, order=order, ignore_index=ignore_index, **labels
+        )
 
     index = get_axis(df, axis)
-    new_index = _assignlevel(index, frame=frame, order=order, **labels)
+    new_index = _assignlevel(
+        index, frame=frame, order=order, ignore_index=ignore_index, **labels
+    )
     return df.set_axis(new_index, axis=axis)
 
 

tests/test_core.py

@@ -81,6 +81,21 @@ def test_assignlevel_index(sidx: Index, midx: MultiIndex):
         ),
     )
 
+    # Check adding from dataframe w/o auto-alignment
+    assert_index_equal(
+        assignlevel(
+            midx,
+            DataFrame(
+                {"new": [1, 2, 3], "new2": 2}, index=Index([2, 1, 3], name="num")
+            ),
+            ignore_index=True,
+        ),
+        MultiIndex.from_arrays(
+            [midx.get_level_values(0), midx.get_level_values(1), [1, 2, 3], [2, 2, 2]],
+            names=["str", "num", "new", "new2"],
+        ),
+    )
+
     # Check wrong types
     with pytest.raises(ValueError):
         assignlevel(midx, "no-frame")