Repository: pyjanitor-devs/pyjanitor

Issue:
Allow pivot_wider and pivot_longer use janitor.select_columns syntax
Using `janitor.columns` syntax will allow for more flexible and dynamic column selection in pivot_wider and pivot_longer. #790 


Discussion:

Resolving patch:
diff --git a/janitor/functions.py b/janitor/functions.py
--- a/janitor/functions.py
+++ b/janitor/functions.py
@@ -44,7 +44,6 @@
     _data_checks_pivot_longer,
     _data_checks_pivot_wider,
     _grid_computation,
-    _pivot_longer_pattern_match,
     _replace_empty_string_with_none,
     _replace_original_empty_string_with_none,
     _select_columns,
@@ -3077,9 +3076,11 @@ def select_columns(
     ::
 
         import pandas as pd
+        import janitor
         import numpy as np
         import datetime
         import re
+        from janitor import patterns
 
         df = pd.DataFrame(
                 {
@@ -3157,6 +3158,15 @@ def select_columns(
         0    4.0     8     E     T     a 2018-01-01
         1    NaN     5   NaN     U     b 2018-01-01
 
+        # same as above, with janitor.patterns
+        # simply a wrapper around re.compile
+
+        df.select_columns(patterns("\\d+"))
+
+           code1 code2 type1 type2 code3      type3
+        0    4.0     8     E     T     a 2018-01-01
+        1    NaN     5   NaN     U     b 2018-01-01
+
     - Select via a list (you can combine any of the previous options)::
 
         df.select_columns(["id", "code*", slice("code", "code2")])
@@ -3216,9 +3226,6 @@ def select_columns(
         search_column_names,
         [str, callable, Pattern, slice, list],
     )
-    if isinstance(search_column_names, list):
-        for label in search_column_names:
-            check("column label", label, [str, slice, Pattern, callable])
 
     full_column_list = _select_columns(search_column_names, df)
 
@@ -5533,15 +5540,15 @@ def pivot_longer(
     :param df: A pandas dataframe.
     :param index: Name(s) of columns to use as identifier variables.
         Should be either a single column name, or a list/tuple of
-        column names. You can also dynamically select column names
-        by using a regular expression with the `janitor.patterns`
-        function. Index should be a list of tuples if the columns
-        are a MultiIndex.
+        column names. The `janitor.select_columns` syntax is supported here,
+        allowing for flexible and dynamic column selection.
+        Index should be a list of tuples if the columns are a MultiIndex.
     :param column_names: Name(s) of columns to unpivot. Should be either
-        a single column name or a list/tuple of column names. You can also
-        dynamically select column names by using a regular expression
-        with the `janitor.patterns` function. Column_names should be a
-        list of tuples if the columns are a MultiIndex.
+        a single column name or a list/tuple of column names.
+        The `janitor.select_columns` syntax is supported here,
+        allowing for flexible and dynamic column selection.
+        Column_names should be a list of tuples
+        if the columns are a MultiIndex.
     :param names_to: Name of new column as a string that will contain
         what were previously the column names in `column_names`.
         The default is `variable` if no value is provided. It can
@@ -5640,10 +5647,6 @@ def pivot_longer(
         ignore_index,
     )
 
-    df, index, column_names = _pivot_longer_pattern_match(
-        df, index, column_names
-    )
-
     df = _computations_pivot_longer(
         df,
         index,
@@ -5824,15 +5827,22 @@ def pivot_wider(
     :param df: A pandas dataframe.
     :param index: Name(s) of columns to use as identifier variables.
         Should be either a single column name, or a list of column names.
+        The `janitor.select_columns` syntax is supported here,
+        allowing for flexible and dynamic column selection.
         If `index` is not provided, the current dataframe's index is used.
     :param names_from: Name(s) of column(s) to use to make the new
         dataframe's columns. Should be either a single column name, or a
-        list of column names. A label or labels must be provided for
-        ``names_from``.
+        list of column names.
+        The `janitor.select_columns` syntax is supported here,
+        allowing for flexible and dynamic column selection.
+        A label or labels must be provided for ``names_from``.
     :param values_from: Name(s) of column(s) that will be used for populating
         the new dataframe's values. Should be either a single column name,
-        or a list of column names. If ``values_from`` is not specified, all
-        remaining columns will be used. If `flatten_levels` is ``False``,
+        or a list of column names.
+        The `janitor.select_columns` syntax is supported here,
+        allowing for flexible and dynamic column selection.
+        If ``values_from`` is not specified,
+        all remaining columns will be used. If `flatten_levels` is ``False``,
         a MultiIndex dataframe is created.
     :param names_sort: Default is `False`. Sorts columns by order of
         appearance. Applicable only if ``flatten_levels`` is `True`.
diff --git a/janitor/utils.py b/janitor/utils.py
--- a/janitor/utils.py
+++ b/janitor/utils.py
@@ -851,37 +851,6 @@ def _data_checks_pivot_longer(
     )
 
 
-def _pivot_longer_pattern_match(
-    df: pd.DataFrame,
-    index: Optional[Pattern] = None,
-    column_names: Optional[Pattern] = None,
-) -> Tuple:
-    """
-    This checks if a pattern (regular expression) is supplied
-    to index or column_names and extracts the column labels
-    that match the given regular expression.
-
-    A dataframe, along with the `index` and `column_names` are
-    returned.
-    """
-
-    # TODO: allow `janitor.patterns` to accept a list/tuple
-    # of regular expresssions.
-    if isinstance(column_names, Pattern):
-        column_names = [
-            column_name
-            for column_name in df
-            if column_names.search(column_name)
-        ]
-
-    if isinstance(index, Pattern):
-        index = [
-            column_name for column_name in df if index.search(column_name)
-        ]
-
-    return df, index, column_names
-
-
 def __tile_compat(arr: pd.Index, df: pd.DataFrame) -> pd.Index:
     """
     Repeats index multiple times.
@@ -1009,7 +978,8 @@ def _pivot_longer_extractions(
     if dot_value:
         # the extra indices (np.arange(len(df))) serve as a way to
         # associate each unpivoted value to a unique row identifier.
-        # It also aids in sorting by appearance. append is used in set_index,
+        # It also aids in sorting by appearance.
+        # append is used in set_index,
         # in case the user wishes to keep the original index.
         if index:
             df = df.set_index([*index, np.arange(len(df))], append=True)
@@ -1207,9 +1177,11 @@ def _computations_pivot_longer(
     """
 
     if index is not None:
+        index = _select_columns(index, df)
         check_column(df, index, present=True)
 
     if column_names is not None:
+        column_names = _select_columns(column_names, df)
         check_column(df, column_names, present=True)
 
     if (
@@ -1350,6 +1322,7 @@ def _data_checks_pivot_wider(
         if isinstance(index, str):
             index = [index]
         check("index", index, [list])
+        index = _select_columns(index, df)
         check_column(df, index, present=True)
 
     if names_from is None:
@@ -1360,14 +1333,13 @@ def _data_checks_pivot_wider(
     if isinstance(names_from, str):
         names_from = [names_from]
     check("names_from", names_from, [list])
+    names_from = _select_columns(names_from, df)
     check_column(df, names_from, present=True)
 
     if values_from is not None:
         check("values_from", values_from, [list, str])
-        if isinstance(values_from, str):
-            check_column(df, [values_from], present=True)
-        else:
-            check_column(df, values_from, present=True)
+        values_from = _select_columns(values_from, df)
+        check_column(df, values_from, present=True)
 
     check("names_sort", names_sort, [bool])
 
@@ -1389,17 +1361,8 @@ def _data_checks_pivot_wider(
     if names_sep is not None:
         check("names_sep", names_sep, [str])
 
-    # cant apply the `check` function here
-    # because of the callable type
-    if aggfunc is not None and not any(
-        (isinstance(aggfunc, (str, list, dict)), callable(aggfunc))
-    ):
-        raise TypeError(
-            """
-            aggfunc should be either a function,
-            string, list, or a dictionary.
-            """
-        )
+    if aggfunc is not None:
+        check("aggfunc", aggfunc, [str, list, dict, callable])
 
     if fill_value is not None:
         check("fill_value", fill_value, [int, float, str])
@@ -1495,6 +1458,9 @@ def _computations_pivot_wider(
 
     df = df.loc[:, values_from]
 
+    if df.shape[-1] == 1:
+        df = df.iloc[:, 0]  # aligns with expected output if pd.pivot is used
+
     aggfunc_index = None
     if aggfunc is not None:
         aggfunc_index = list(range(df.index.nlevels))
@@ -1690,7 +1656,12 @@ def _computations_as_categorical(df: pd.DataFrame, **kwargs) -> pd.DataFrame:
 
 
 @functools.singledispatch
-def _select_columns(columns_to_select: str, df):
+def _select_columns(columns_to_select, df):
+    raise TypeError("This type is not supported in column selection.")
+
+
+@_select_columns.register(str)  # noqa: F811
+def _column_sel_dispatch(columns_to_select, df):  # noqa: F811
     """
     Base function for column selection.
     Applies only to strings.
@@ -1704,7 +1675,7 @@ def _select_columns(columns_to_select: str, df):
     elif columns_to_select in df.columns:
         filtered_columns = [columns_to_select]
     if not filtered_columns:
-        raise NameError(f"No match was returned for '{columns_to_select}'")
+        raise KeyError(f"No match was returned for '{columns_to_select}'")
     return filtered_columns
 
 
@@ -1842,11 +1813,23 @@ def _column_sel_dispatch(columns_to_select, df):  # noqa: F811
     ]
 
     if not filtered_columns:
-        raise ValueError("No column name matched the regular expression.")
+        raise KeyError("No column name matched the regular expression.")
 
     return filtered_columns
 
 
+@_select_columns.register(tuple)  # noqa: F811
+def _column_sel_dispatch(columns_to_select, df):  # noqa: F811
+    """
+    Base function for column selection.
+    Applies only to tuple type, and applies to MultiIndex columns.
+    A list/pd.Index of column names is returned.
+    """
+    if columns_to_select not in df.columns:
+        raise KeyError(f"No match was returned for {columns_to_select}")
+    return [columns_to_select]
+
+
 @_select_columns.register(list)  # noqa: F811
 def _column_sel_dispatch(columns_to_select, df):  # noqa: F811
     """
@@ -1856,6 +1839,7 @@ def _column_sel_dispatch(columns_to_select, df):  # noqa: F811
     or a combination of these types.
     A list/pd.Index of column names is returned.
     """
+
     filtered_columns = []
 
     for entry in columns_to_select:


Test patch:
diff --git a/tests/functions/test_pivot_longer.py b/tests/functions/test_pivot_longer.py
--- a/tests/functions/test_pivot_longer.py
+++ b/tests/functions/test_pivot_longer.py
@@ -92,7 +92,7 @@ def df_multi():
 column_labels = [{"region": 2007}, {"2007", "2009"}]
 names_to_labels = [1, {12, "newnames"}]
 
-index_does_not_exist = ["Region", [2007, "region"]]
+index_does_not_exist = ["Region", ["207", "region"]]
 column_does_not_exist = ["two thousand and seven", ("2007", 2009)]
 
 index_type_checks = (
@@ -169,15 +169,15 @@ def test_subtype_names_to(df, names_to):
 
 @pytest.mark.parametrize("df,index", index_presence_checks)
 def test_presence_index(df, index):
-    """Raise ValueError if labels in `index` do not exist."""
-    with pytest.raises(ValueError):
+    """Raise KeyError if labels in `index` do not exist."""
+    with pytest.raises(KeyError):
         df.pivot_longer(index=index)
 
 
 @pytest.mark.parametrize("df,column", column_presence_checks)
 def test_presence_columns(df, column):
-    """Raise ValueError if labels in `column_names` do not exist."""
-    with pytest.raises(ValueError):
+    """Raise KeyError if labels in `column_names` do not exist."""
+    with pytest.raises(KeyError):
         df.pivot_longer(column_names=column)
 
 
diff --git a/tests/functions/test_pivot_wider.py b/tests/functions/test_pivot_wider.py
--- a/tests/functions/test_pivot_wider.py
+++ b/tests/functions/test_pivot_wider.py
@@ -216,28 +216,28 @@ def test_names_from_None(df_checks_output):
 
 
 def test_presence_index1(df_checks_output):
-    """Raise ValueError if labels in `index` do not exist."""
-    with pytest.raises(ValueError):
+    """Raise KeyError if labels in `index` do not exist."""
+    with pytest.raises(KeyError):
         df_checks_output.pivot_wider(index="geo", names_from="variable")
 
 
 def test_presence_index2(df_checks_output):
-    """Raise ValueError if labels in `index` do not exist."""
-    with pytest.raises(ValueError):
+    """Raise KeyError if labels in `index` do not exist."""
+    with pytest.raises(KeyError):
         df_checks_output.pivot_wider(
             index=["geoid", "Name"], names_from="variable"
         )
 
 
 def test_presence_names_from1(df_checks_output):
-    """Raise ValueError if labels in `names_from` do not exist."""
-    with pytest.raises(ValueError):
+    """Raise KeyError if labels in `names_from` do not exist."""
+    with pytest.raises(KeyError):
         df_checks_output.pivot_wider(index="geoid", names_from="estmt")
 
 
 def test_presence_names_from2(df_checks_output):
-    """Raise ValueError if labels in `names_from` do not exist."""
-    with pytest.raises(ValueError):
+    """Raise KeyError if labels in `names_from` do not exist."""
+    with pytest.raises(KeyError):
         df_checks_output.pivot_wider(index="geoid", names_from=["estimat"])
 
 
@@ -546,7 +546,7 @@ def test_index_names_sort_True():
     assert_frame_equal(result, expected_output)
 
 
-@pytest.fixture()
+@pytest.fixture
 def df_aggfunc():
     return pd.DataFrame(
         [
@@ -572,6 +572,49 @@ def df_aggfunc():
     )
 
 
+@pytest.fixture
+def df_aggfunc_multiple_names_from():
+    return pd.DataFrame(
+        {
+            "A": [
+                "foo",
+                "foo",
+                "foo",
+                "foo",
+                "foo",
+                "bar",
+                "bar",
+                "bar",
+                "bar",
+            ],
+            "B": [
+                "one",
+                "one",
+                "one",
+                "two",
+                "two",
+                "one",
+                "one",
+                "two",
+                "two",
+            ],
+            "C": [
+                "small",
+                "large",
+                "large",
+                "small",
+                "small",
+                "large",
+                "small",
+                "small",
+                "large",
+            ],
+            "D": [1, 2, 2, 3, 3, 4, 5, 6, 7],
+            "E": [2, 4, 5, 5, 6, 6, 8, 9, 9],
+        }
+    )
+
+
 def test_aggfunc(df_aggfunc):
     """Test output when `aggfunc` is provided."""
     expected = pd.DataFrame(
@@ -610,3 +653,64 @@ def test_aggfunc_list(df_aggfunc):
         index="V4", names_from="variable", aggfunc=["sum"], flatten_levels=True
     )
     assert_frame_equal(result, expected)
+
+
+def test_aggfunc_multiple_names_from(df_aggfunc_multiple_names_from):
+    """Test output when ``names_from`` is more than one."""
+    expected = pd.DataFrame(
+        {
+            "A": ["bar", "bar", "foo", "foo"],
+            "C": ["large", "small", "large", "small"],
+            "one_D": [4.0, 5.0, 2.0, 1.0],
+            "two_D": [7.0, 6.0, np.nan, 3.0],
+            "one_E": [6.0, 8.0, 4.5, 2.0],
+            "two_E": [9.0, 9.0, np.nan, 5.5],
+        }
+    )
+    result = df_aggfunc_multiple_names_from.pivot_wider(
+        index=["A", "C"],
+        names_from=["B"],
+        values_from=["D", "E"],
+        aggfunc=["mean"],
+        flatten_levels=True,
+    )
+    assert_frame_equal(result, expected)
+
+
+def test_df_multiple_aggfuncs():
+    """Test output when ``aggfunc`` is more than one."""
+
+    df_frame = pd.DataFrame(
+        [
+            {"A": "foo", "B": "one", "C": "small", "D": 1, "E": 2},
+            {"A": "foo", "B": "one", "C": "large", "D": 2, "E": 4},
+            {"A": "foo", "B": "one", "C": "large", "D": 2, "E": 5},
+            {"A": "foo", "B": "one", "C": "small", "D": 3, "E": 5},
+            {"A": "foo", "B": "one", "C": "small", "D": 3, "E": 6},
+            {"A": "bar", "B": "one", "C": "large", "D": 4, "E": 6},
+            {"A": "bar", "B": "one", "C": "small", "D": 5, "E": 8},
+            {"A": "bar", "B": "one", "C": "small", "D": 6, "E": 9},
+            {"A": "bar", "B": "one", "C": "large", "D": 7, "E": 9},
+        ]
+    )
+
+    expected = pd.DataFrame(
+        {
+            "A": ["bar", "bar", "foo", "foo"],
+            "C": ["large", "small", "large", "small"],
+            "one_D_mean": [5.5, 5.5, 2.0, 2.3333333333333335],
+            "one_D_sum": [11, 11, 4, 7],
+            "one_E_mean": [7.5, 8.5, 4.5, 4.333333333333333],
+            "one_E_sum": [15, 17, 9, 13],
+        }
+    )
+
+    result = df_frame.pivot_wider(
+        index=["A", "C"],
+        names_from=["B"],
+        values_from=["D", "E"],
+        aggfunc=["mean", "sum"],
+        flatten_levels=True,
+    )
+
+    assert_frame_equal(result, expected)
diff --git a/tests/functions/test_select_columns.py b/tests/functions/test_select_columns.py
--- a/tests/functions/test_select_columns.py
+++ b/tests/functions/test_select_columns.py
@@ -44,7 +44,7 @@ def test_select_columns_glob_inputs(dataframe, invert, expected):
 )
 def test_select_columns_missing_columns(dataframe, columns):
     """Check that passing non-existent column names or search strings raises NameError"""  # noqa: E501
-    with pytest.raises(NameError):
+    with pytest.raises(KeyError):
         dataframe.select_columns(search_column_names=columns)
 
 
diff --git a/tests/utils/test__select_column.py b/tests/utils/test__select_column.py
--- a/tests/utils/test__select_column.py
+++ b/tests/utils/test__select_column.py
@@ -47,6 +47,19 @@ def df1():
     )
 
 
+@pytest.fixture
+def df_tuple():
+    frame = pd.DataFrame(
+        {
+            "A": {0: "a", 1: "b", 2: "c"},
+            "B": {0: 1, 1: 3, 2: 5},
+            "C": {0: 2, 1: 4, 2: 6},
+        }
+    )
+    frame.columns = [list("ABC"), list("DEF")]
+    return frame
+
+
 def test_type(df):
     """Raise TypeError if `columns_to_select` is the wrong type."""
     with pytest.raises(TypeError):
@@ -59,12 +72,12 @@ def test_type(df):
 
 def test_strings_do_not_exist(df):
     """
-    Raise NameError if `columns_to_select` is a string
+    Raise KeyError if `columns_to_select` is a string
     and does not exist in the dataframe's columns.
     """
-    with pytest.raises(NameError):
+    with pytest.raises(KeyError):
         _select_columns("word", df)
-    with pytest.raises(NameError):
+    with pytest.raises(KeyError):
         _select_columns("*starter", df)
 
 
@@ -128,13 +141,22 @@ def test_callable_no_match(df):
 
 def test_regex_presence(df):
     """
-    Raise ValueError if `columns_to_select` is a regex
+    Raise KeyError if `columns_to_select` is a regex
     and none of the column names match.
     """
-    with pytest.raises(ValueError):
+    with pytest.raises(KeyError):
         _select_columns(re.compile(r"^\d+"), df)
 
 
+def test_tuple_presence(df_tuple):
+    """
+    Raise KeyError if `columns_to_select` is a tuple
+    and no match is returned.
+    """
+    with pytest.raises(KeyError):
+        _select_columns(("A", "C"), df_tuple)
+
+
 def test_strings(df1):
     """Test _select_columns function on strings."""
     assert _select_columns("id", df1) == ["id"]
@@ -250,6 +272,13 @@ def test_regex(df1):
     )
 
 
+def test_tuple(df_tuple):
+    """Test _select_columns function on tuple."""
+    assert _select_columns(("A", "D"), df_tuple) == list(
+        df_tuple.loc[:, [("A", "D")]].columns
+    )
+
+
 def test_list_various(df1):
     """Test _select_columns function on list type."""