Merge pull request #266 from rs-station/pandas_signatures

Update `rs.DataSet.reset_index()` call signature to matches `pandas >v1.5`

reciprocalspaceship/dataset.py

@@ -258,7 +258,14 @@ def set_index(
         )
 
     def reset_index(
-        self, level=None, drop=False, inplace=False, col_level=0, col_fill=""
+        self,
+        level=None,
+        drop=False,
+        inplace=False,
+        col_level=0,
+        col_fill="",
+        allow_duplicates=lib.no_default,
+        names=None,
     ):
         """
         Reset the index or a specific level of a MultiIndex.
@@ -281,6 +288,12 @@ def reset_index(
         col_fill : object
             If the columns have multiple levels, determines how the other
             levels are named. If None then the index name is repeated.
+        allow_duplicates : bool
+            Allow duplicate column labels to be created.
+        names : int, str, tuple, list
+            Using the given string, rename the DataSet column which contains the
+            index data. If the DataSet has a MultiIndex, this has to be a list or
+            tuple with length equal to the number of levels.
 
         Returns
         -------
@@ -317,6 +330,8 @@ def _handle_cached_dtypes(dataset, columns, drop):
                 inplace=inplace,
                 col_level=col_level,
                 col_fill=col_fill,
+                allow_duplicates=allow_duplicates,
+                names=names,
             )
             _handle_cached_dtypes(self, columns, drop)
             return
@@ -327,6 +342,8 @@ def _handle_cached_dtypes(dataset, columns, drop):
                 inplace=inplace,
                 col_level=col_level,
                 col_fill=col_fill,
+                allow_duplicates=allow_duplicates,
+                names=names,
             )
             dataset._index_dtypes = dataset._index_dtypes.copy()
             dataset = _handle_cached_dtypes(dataset, columns, drop)

tests/test_dataset_signatures.py

@@ -0,0 +1,53 @@
+from inspect import signature
+
+import pandas as pd
+import pytest
+from pandas.testing import assert_frame_equal
+
+import reciprocalspaceship as rs
+
+
+def test_reset_index_dataseries():
+    """
+    Minimal example from GH#223
+    """
+    result = rs.DataSeries(range(10)).reset_index()
+    expected = pd.Series(range(10)).reset_index()
+    expected = rs.DataSet(expected)
+    assert_frame_equal(result, expected)
+
+
+def test_reset_index_signature(dataset_hkl):
+    """
+    Test call signature of rs.DataSet.reset_index() matches call signature of
+    pd.DataFrame.reset_index() using default parameters
+    """
+    df = pd.DataFrame(dataset_hkl)
+    sig = signature(pd.DataFrame.reset_index)
+    bsig = sig.bind(df)
+    bsig.apply_defaults()
+
+    expected = df.reset_index(*bsig.args[1:], **bsig.kwargs)
+    result = dataset_hkl.reset_index(*bsig.args[1:], **bsig.kwargs)
+    result = pd.DataFrame(result)
+
+    assert_frame_equal(result, expected)
+
+
+@pytest.mark.parametrize("names", ["H", "K", ["H", "K"]])
+def test_set_index_signature(dataset_hkl, names):
+    """
+    Test call signature of rs.DataSet.set_index() matches call signature of
+    pd.DataFrame.set_index() using default parameters
+    """
+    ds = dataset_hkl.reset_index()
+    df = pd.DataFrame(ds)
+    sig = signature(pd.DataFrame.set_index)
+    bsig = sig.bind(df, names)
+    bsig.apply_defaults()
+
+    expected = df.set_index(*bsig.args[1:], **bsig.kwargs)
+    result = ds.set_index(*bsig.args[1:], **bsig.kwargs)
+    result = pd.DataFrame(result)
+
+    assert_frame_equal(result, expected)