Repository: pyjanitor-devs/pyjanitor

Issue:
extend `row_to_names` to support multiindex
allow `row_to_names` to convert multiple rows to multiindex


Discussion:

Resolving patch:
diff --git a/janitor/functions/row_to_names.py b/janitor/functions/row_to_names.py
--- a/janitor/functions/row_to_names.py
+++ b/janitor/functions/row_to_names.py
@@ -2,23 +2,25 @@
 
 import warnings
 
+import numpy as np
 import pandas as pd
 import pandas_flavor as pf
 
-from janitor.utils import check
+from janitor.utils import check, deprecated_alias
 
 
 @pf.register_dataframe_method
+@deprecated_alias(row_number="row_numbers", remove_row="remove_rows")
 def row_to_names(
     df: pd.DataFrame,
-    row_number: int = 0,
-    remove_row: bool = False,
+    row_numbers: int = 0,
+    remove_rows: bool = False,
     remove_rows_above: bool = False,
     reset_index: bool = False,
 ) -> pd.DataFrame:
-    """Elevates a row to be the column names of a DataFrame.
+    """Elevates a row, or rows, to be the column names of a DataFrame.
 
-    This method mutates the original DataFrame.
+    This method does not mutate the original DataFrame.
 
     Contains options to remove the elevated row from the DataFrame along with
     removing the rows above the selected row.
@@ -37,10 +39,14 @@ def row_to_names(
         0  nums  chars
         1     6      x
         2     9      y
-        >>> df.row_to_names(0, remove_row=True, reset_index=True)
+        >>> df.row_to_names(0, remove_rows=True, reset_index=True)
           nums chars
         0    6     x
         1    9     y
+        >>> df.row_to_names([0,1], remove_rows=True, reset_index=True)
+          nums chars
+          6    x
+        0    9     y
 
         Remove rows above the elevated row and the elevated row itself.
 
@@ -56,25 +62,32 @@ def row_to_names(
         1  nums  chars
         2     6      x
         3     9      y
-        >>> df.row_to_names(1, remove_row=True, remove_rows_above=True, reset_index=True)
+        >>> df.row_to_names(1, remove_rows=True, remove_rows_above=True, reset_index=True)
           nums chars
         0    6     x
         1    9     y
 
     Args:
         df: A pandas DataFrame.
-        row_number: Position of the row containing the variable names.
-            Note that indexing starts from 0. Defaults to 0 (first row).
-        remove_row: Whether the row should be removed from the DataFrame.
-        remove_rows_above: Whether the rows above the selected row should
+        row_numbers: Position of the row(s) containing the variable names.
+            Note that indexing starts from 0. It can also be a list,
+            in which case, a MultiIndex column is created.
+            Defaults to 0 (first row).
+        remove_row: Whether the row(s) should be removed from the DataFrame.
+        remove_rows_above: Whether the row(s) above the selected row should
             be removed from the DataFrame.
         reset_index: Whether the index should be reset on the returning DataFrame.
 
     Returns:
         A pandas DataFrame with set column names.
     """  # noqa: E501
+    if not pd.options.mode.copy_on_write:
+        df = df.copy()
 
-    check("row_number", row_number, [int])
+    check("row_number", row_numbers, [int, list])
+    if isinstance(row_numbers, list):
+        for entry in row_numbers:
+            check("entry in the row_number argument", entry, [int])
 
     warnings.warn(
         "The function row_to_names will, in the official 1.0 release, "
@@ -82,17 +95,39 @@ def row_to_names(
         "You can prepare for this change right now by explicitly setting "
         "`reset_index=True` when calling on `row_to_names`."
     )
-
-    df.columns = df.iloc[row_number, :]
+    # should raise if positional indexers are missing
+    # IndexError: positional indexers are out-of-bounds
+    headers = df.iloc[row_numbers]
+    if isinstance(headers, pd.DataFrame) and (len(headers) == 1):
+        headers = headers.squeeze()
+    if isinstance(headers, pd.Series):
+        headers = pd.Index(headers)
+    else:
+        headers = [entry.array for _, entry in headers.items()]
+        headers = pd.MultiIndex.from_tuples(headers)
+
+    df.columns = headers
     df.columns.name = None
 
-    if remove_row:
-        df = df.drop(df.index[row_number])
-
+    df_index = df.index
     if remove_rows_above:
-        df = df.drop(df.index[range(row_number)])
-
+        if isinstance(row_numbers, list):
+            if not (np.diff(row_numbers) == 1).all():
+                raise ValueError(
+                    "The remove_rows_above argument is applicable "
+                    "only if the row_numbers argument is an integer, "
+                    "or the integers in a list are consecutive increasing, "
+                    "with a difference of 1."
+                )
+            tail = row_numbers[0]
+        else:
+            tail = row_numbers
+        df = df.iloc[tail:]
+    if remove_rows:
+        if isinstance(row_numbers, int):
+            row_numbers = [row_numbers]
+        df_index = df.index.symmetric_difference(df_index[row_numbers])
+        df = df.loc[df_index]
     if reset_index:
-        df = df.reset_index(drop=["index"])
-
+        df.index = range(len(df))
     return df


Test patch:
diff --git a/tests/functions/test_row_to_names.py b/tests/functions/test_row_to_names.py
--- a/tests/functions/test_row_to_names.py
+++ b/tests/functions/test_row_to_names.py
@@ -12,9 +12,31 @@ def test_row_to_names(dataframe):
     assert df.columns[4] == "Basel"
 
 
+@pytest.mark.functions
+def test_row_to_names_single_list(dataframe):
+    "Test output if row_numbers is a list, and contains a single item."
+    df = dataframe.row_to_names([2])
+    assert df.columns[0] == 3
+    assert df.columns[1] == 3.234_612_5
+    assert df.columns[2] == 3
+    assert df.columns[3] == "lion"
+    assert df.columns[4] == "Basel"
+
+
+@pytest.mark.functions
+def test_row_to_names_list(dataframe):
+    "Test output if row_numbers is a list."
+    df = dataframe.row_to_names([1, 2])
+    assert df.columns[0] == (2, 3)
+    assert df.columns[1] == (2.456_234, 3.234_612_5)
+    assert df.columns[2] == (2, 3)
+    assert df.columns[3] == ("leopard", "lion")
+    assert df.columns[4] == ("Shanghai", "Basel")
+
+
 @pytest.mark.functions
 def test_row_to_names_delete_this_row(dataframe):
-    df = dataframe.row_to_names(2, remove_row=True)
+    df = dataframe.row_to_names(2, remove_rows=True)
     assert df.iloc[2, 0] == 1
     assert df.iloc[2, 1] == 1.234_523_45
     assert df.iloc[2, 2] == 1
@@ -26,7 +48,7 @@ def test_row_to_names_delete_this_row(dataframe):
 def test_row_to_names_delete_the_row_without_resetting_index(dataframe):
     """Test that executes row_to_names while deleting the given row
     index while not resetting the index"""
-    df = dataframe.row_to_names(2, remove_row=True)
+    df = dataframe.row_to_names(2, remove_rows=True)
     expected_index = pd.Index([0, 1, 3, 4, 5, 6, 7, 8])
     pd.testing.assert_index_equal(df.index, expected_index)
 
@@ -41,6 +63,28 @@ def test_row_to_names_delete_above(dataframe):
     assert df.iloc[0, 4] == "Basel"
 
 
+@pytest.mark.functions
+def test_row_to_names_delete_above_list(dataframe):
+    "Test output if row_numbers is a list"
+    df = dataframe.row_to_names([2, 3], remove_rows_above=True)
+    assert df.iloc[0, 0] == 3
+    assert df.iloc[0, 1] == 3.234_612_5
+    assert df.iloc[0, 2] == 3
+    assert df.iloc[0, 3] == "lion"
+    assert df.iloc[0, 4] == "Basel"
+
+
+@pytest.mark.functions
+def test_row_to_names_delete_above_list_non_consecutive(dataframe):
+    "Raise if row_numbers is a list, but non consecutive"
+    msg = "The remove_rows_above argument is applicable "
+    msg += "only if the row_numbers argument is an integer, "
+    msg += "or the integers in a list are consecutive increasing, "
+    msg += "with a difference of 1."
+    with pytest.raises(ValueError, match=msg):
+        dataframe.row_to_names([1, 3], remove_rows_above=True)
+
+
 @pytest.mark.functions
 def test_row_to_names_delete_above_without_resetting_index(dataframe):
     """Test that executes row_to_names while deleting the all rows
@@ -63,6 +107,6 @@ def test_row_to_names_delete_above_with_resetting_index(dataframe):
 def test_row_to_names_delete_the_row_with_resetting_index(dataframe):
     """Test that executes row_to_names while deleting the given row
     index while resetting the index"""
-    df = dataframe.row_to_names(2, remove_row=True, reset_index=True)
+    df = dataframe.row_to_names(2, remove_rows=True, reset_index=True)
     expected_index = pd.RangeIndex(start=0, stop=8, step=1)
     pd.testing.assert_index_equal(df.index, expected_index)