Repository: pyjanitor-devs/pyjanitor

Issue:
Suggestions for `select_columns`
- [x] Allow selection on a level (for MultiIndex columns)
- [x] Return pd.Index if possible (`loc` can handle column selection as an Index)
- [x] Extend slicing to MultiIndex columns via tuples (currently handles only string/integers)
- [x] Allow selection to be generic to hashable, not just strings/tuples (columns can be numeric dtypes)
- [x] As much as possible, pass column presence/absence check to `.loc` (which raises a KeyError if label is not found)


Discussion:

Resolving patch:
diff --git a/janitor/functions/select_columns.py b/janitor/functions/select_columns.py
--- a/janitor/functions/select_columns.py
+++ b/janitor/functions/select_columns.py
@@ -1,11 +1,11 @@
 """Implementation of select_columns"""
+from typing import Optional, Union
 import pandas_flavor as pf
 import pandas as pd
-
-from janitor.utils import deprecated_alias
+from pandas.api.types import is_list_like
+from janitor.utils import deprecated_alias, check
 
 from janitor.functions.utils import _select_column_names
-from pandas.api.types import is_list_like
 
 
 @pf.register_dataframe_method
@@ -13,6 +13,7 @@
 def select_columns(
     df: pd.DataFrame,
     *args,
+    level: Optional[Union[int, str]] = None,
     invert: bool = False,
 ) -> pd.DataFrame:
     """
@@ -48,6 +49,8 @@ def select_columns(
         a callable which is applicable to each Series in the DataFrame,
         or variable arguments of all the aforementioned.
         A sequence of booleans is also acceptable.
+    :param level: Determines which level in the columns should be used for the
+        column selection.
     :param invert: Whether or not to invert the selection.
         This will result in the selection of the complement of the columns
         provided.
@@ -62,8 +65,26 @@ def select_columns(
             search_column_names.extend(arg)
         else:
             search_column_names.append(arg)
-    full_column_list = _select_column_names(search_column_names, df)
-
+    if level is not None:
+        # goal here is to capture the original columns
+        # trim the df.columns to the specified level only,
+        # and apply the selection (_select_column_names)
+        # to get the relevant column labels.
+        # note that no level is dropped; if there are three levels,
+        # then three levels are returned, with the specified labels
+        # selected/deselected.
+        # A copy of the dataframe is made via set_axis,
+        # to avoid mutating the original dataframe.
+        df_columns = df.columns
+        check("level", level, [int, str])
+        full_column_list = df_columns.get_level_values(level)
+        full_column_list = _select_column_names(
+            search_column_names, df.set_axis(full_column_list, axis=1)
+        )
+        full_column_list = df_columns.isin(full_column_list, level=level)
+        full_column_list = df_columns[full_column_list]
+    else:
+        full_column_list = _select_column_names(search_column_names, df)
     if invert:
         return df.drop(columns=full_column_list)
     return df.loc[:, full_column_list]
diff --git a/janitor/functions/utils.py b/janitor/functions/utils.py
--- a/janitor/functions/utils.py
+++ b/janitor/functions/utils.py
@@ -12,6 +12,9 @@
     union_categoricals,
     is_scalar,
     is_list_like,
+    is_datetime64_dtype,
+    is_string_dtype,
+    is_categorical_dtype,
 )
 import numpy as np
 from multipledispatch import dispatch
@@ -212,7 +215,16 @@ def _select_column_names(columns_to_select, df):
     """
     if columns_to_select in df.columns:
         return [columns_to_select]
-    raise KeyError(f"No match was returned for '{columns_to_select}'.")
+    raise KeyError(f"No match was returned for {columns_to_select}.")
+
+
+def _is_str_or_cat(df_columns):
+    """Check if the column is a string or categorical with strings."""
+    if is_string_dtype(df_columns):
+        return True
+    if is_categorical_dtype(df_columns):
+        return is_string_dtype(df_columns.categories)
+    return False
 
 
 @_select_column_names.register(str)  # noqa: F811
@@ -221,32 +233,26 @@ def _column_sel_dispatch(columns_to_select, df):  # noqa: F811
     Base function for column selection.
     Applies only to strings.
     It is also applicable to shell-like glob strings,
-    specifically, the `*`.
+    which are supported by `fnmatch`.
     A list/pandas Index of matching column names is returned.
     """
     df_columns = df.columns
-    if pd.api.types.is_string_dtype(df_columns):
-        if (
-            "*" in columns_to_select
-        ):  # shell-style glob string (e.g., `*_thing_*`)
-            return fnmatch.filter(df_columns, columns_to_select)
+
+    if _is_str_or_cat(df_columns):
         if columns_to_select in df_columns:
             return [columns_to_select]
-        raise KeyError(f"No match was returned for '{columns_to_select}'.")
-    if pd.api.types.is_datetime64_any_dtype(df_columns):
-        if not df_columns.is_monotonic_increasing:
-            raise ValueError(
-                "The column is a DatetimeIndex and should be "
-                "monotonic increasing."
-            )
+        outcome = fnmatch.filter(df_columns, columns_to_select)
+        if not outcome:
+            raise KeyError(f"No match was returned for '{columns_to_select}'.")
+        return outcome
+
+    if is_datetime64_dtype(df_columns):
         timestamp = df_columns.get_loc(columns_to_select)
-        if isinstance(timestamp, slice):
+        if not isinstance(timestamp, int):
             return df_columns[timestamp]
         return [df_columns[timestamp]]
-    raise KeyError(
-        f"String('{columns_to_select}') can be applied "
-        "only to string/datetime columns."
-    )
+
+    raise KeyError(f"No match was returned for '{columns_to_select}'.")
 
 
 @_select_column_names.register(re.Pattern)  # noqa: F811
@@ -257,15 +263,16 @@ def _column_sel_dispatch(columns_to_select, df):  # noqa: F811
     `re.compile` is required for the regular expression.
     A pandas Index of matching column names is returned.
     """
-    if pd.api.types.is_string_dtype(df.columns):
-        bools = df.columns.str.contains(
+    df_columns = df.columns
+
+    if _is_str_or_cat(df_columns):
+        bools = df_columns.str.contains(
             columns_to_select, na=False, regex=True
         )
-        return df.columns[bools]
-    raise KeyError(
-        f"Regular expressions('{columns_to_select}') "
-        "can be applied only to string columns."
-    )
+        if not bools.any():
+            raise KeyError(f"No match was returned for {columns_to_select}.")
+        return df_columns[bools]
+    raise KeyError(f"No match was returned for {columns_to_select}.")
 
 
 @_select_column_names.register(slice)  # noqa: F811
@@ -291,13 +298,12 @@ def _column_sel_dispatch(columns_to_select, df):  # noqa: F811
     step_check = None
     method = None
 
-    if not df_columns.is_unique:
+    if not df_columns.is_unique and not df_columns.is_monotonic_increasing:
         raise ValueError(
-            "The column labels are not unique. "
-            "Kindly ensure the labels are unique "
-            "to ensure the correct output."
+            "Non-unique column labels should be monotonic increasing."
         )
-    is_date_column = pd.api.types.is_datetime64_any_dtype(df_columns)
+
+    is_date_column = is_datetime64_dtype(df_columns)
     if is_date_column:
         if not df_columns.is_monotonic_increasing:
             raise ValueError(
@@ -377,6 +383,8 @@ def _column_sel_dispatch(columns_to_select, df):  # noqa: F811
         raise TypeError(
             "The output of the applied callable should be a boolean array."
         )
+    if not filtered_columns.any():
+        raise KeyError(f"No match was returned for {columns_to_select}.")
 
     return df.columns[filtered_columns]
 


Test patch:
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
@@ -1,5 +1,6 @@
 import pandas as pd
 import pytest
+import re
 from pandas.testing import assert_frame_equal
 
 
@@ -51,36 +52,6 @@ def test_select_column_names_missing_columns(dataframe, columns):
         dataframe.select_columns(columns)
 
 
-@pytest.mark.functions
-@pytest.mark.parametrize(
-    "columns",
-    [
-        pytest.param(
-            "a",
-            marks=pytest.mark.xfail(
-                reason="`select_columns` now accepts strings"
-            ),
-        ),
-        pytest.param(
-            ("a", "Bell__Chart"),
-            marks=pytest.mark.xfail(
-                reason="`select_columns` converts list-like into lists"
-            ),
-        ),
-        pytest.param(
-            {"a", "Bell__Chart"},
-            marks=pytest.mark.xfail(
-                reason="`select_columns` converts list-like into lists"
-            ),
-        ),
-    ],
-)
-def test_select_column_names_input(dataframe, columns):
-    """Check that passing an iterable that is not a list raises TypeError."""
-    with pytest.raises(TypeError):
-        dataframe.select_columns(columns)
-
-
 @pytest.mark.functions
 @pytest.mark.parametrize(
     "invert,expected",
@@ -116,20 +87,60 @@ def columns(x):
     assert_frame_equal(df, dataframe[expected])
 
 
-@pytest.mark.xfail(reason="Allow tuples which are acceptable in MultiIndex.")
-def test_MultiIndex():
-    """
-    Raise ValueError if columns is a MultiIndex.
-    """
-    df = pd.DataFrame(
+@pytest.fixture
+def df_tuple():
+    "pytest fixture."
+    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_multiindex(df_tuple):
+    """
+    Test output for a MultiIndex and tuple passed.
+    """
+    assert_frame_equal(
+        df_tuple.select_columns(("A", "D")), df_tuple.loc[:, [("A", "D")]]
+    )
 
-    df.columns = [list("ABC"), list("DEF")]
 
-    with pytest.raises(ValueError):
-        df.select_columns("A")
+def test_level_callable(df_tuple):
+    """
+    Test output if level is supplied for a callable.
+    """
+    expected = df_tuple.select_columns(
+        lambda df: df.name.startswith("A"), level=0
+    )
+    actual = df_tuple.xs("A", axis=1, drop_level=False, level=0)
+    assert_frame_equal(actual, expected)
+
+
+def test_level_regex(df_tuple):
+    """
+    Test output if level is supplied for a regex
+    """
+    expected = df_tuple.select_columns(re.compile("D"), level=1)
+    actual = df_tuple.xs("D", axis=1, drop_level=False, level=1)
+    assert_frame_equal(actual, expected)
+
+
+def test_level_slice(df_tuple):
+    """
+    Test output if level is supplied for a slice
+    """
+    expected = df_tuple.select_columns(slice("F", "D"), level=1)
+    assert_frame_equal(df_tuple, expected)
+
+
+def test_level_str(df_tuple):
+    """
+    Test output if level is supplied for a string.
+    """
+    expected = df_tuple.select_columns("A", level=0, invert=True)
+    assert_frame_equal(df_tuple.drop(columns="A", axis=1, level=0), expected)
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
@@ -103,7 +103,7 @@ def test_col_not_found3(df_dates):
 
 def test_col_not_found4(df_numbers):
     """Raise KeyError if `columns_to_select` is not in df.columns."""
-    with pytest.raises(KeyError, match=r"String\(.+\) can be applied.+"):
+    with pytest.raises(KeyError, match=r"No match was returned.+"):
         _select_column_names("id", df_numbers)
 
 
@@ -123,6 +123,18 @@ def test_strings(df1):
     ]
 
 
+def test_strings_cat(df1):
+    """Test output on categorical columns"""
+    df1.columns = df1.columns.astype("category")
+    assert _select_column_names("id", df1) == ["id"]
+    assert _select_column_names("*type*", df1) == [
+        "type",
+        "type1",
+        "type2",
+        "type3",
+    ]
+
+
 def test_strings_do_not_exist(df):
     """
     Raise KeyError if `columns_to_select` is a string
@@ -151,14 +163,12 @@ def test_strings_dates_range(df_dates):
 
 
 def test_unsorted_dates(df_dates):
-    """Raise Error if the dates are unsorted."""
+    """Test output if the dates are unsorted, and a string is passed."""
     df_dates = df_dates.iloc[:, [10, 4, 7, 2, 1, 3, 5, 6, 8, 9, 11, 0]]
-    with pytest.raises(
-        ValueError,
-        match="The column is a DatetimeIndex and should be "
-        "monotonic increasing.",
-    ):
-        _select_column_names("2011-01-31", df_dates)
+    assert_index_equal(
+        df_dates.loc[:, ["2011-01-31"]].columns,
+        _select_column_names("2011-01-31", df_dates),
+    )
 
 
 def test_regex(df1):
@@ -169,6 +179,15 @@ def test_regex(df1):
     )
 
 
+def test_regex_cat(df1):
+    """Test output on categorical columns"""
+    df1.columns = df1.columns.astype("category")
+    assert_index_equal(
+        _select_column_names(re.compile(r"\d$"), df1),
+        df1.filter(regex=r"\d$").columns,
+    )
+
+
 def test_patterns_warning(df1):
     """
     Check that warning is raised if `janitor.patterns` is used.
@@ -180,14 +199,21 @@ def test_patterns_warning(df1):
         )
 
 
+def test_regex_presence_string_column(df):
+    """
+    Raise KeyError if `columns_to_select` is a regex
+    and does not exist in the dataframe's columns.
+    """
+    with pytest.raises(KeyError, match="No match was returned for.+"):
+        _select_column_names(re.compile("word"), df)
+
+
 def test_regex_presence(df_dates):
     """
     Raise KeyError if `columns_to_select` is a regex
     and the columns is not a string column.
     """
-    with pytest.raises(
-        KeyError, match=r"Regular expressions\(.+\) can be applied.+"
-    ):
+    with pytest.raises(KeyError, match=r"No match was returned.+"):
         _select_column_names(re.compile(r"^\d+"), df_dates)
 
 
@@ -195,8 +221,11 @@ def test_slice_unique():
     """
     Raise ValueError if the columns are not unique.
     """
-    not_unique = pd.DataFrame([], columns=["code", "code", "code1", "code2"])
-    with pytest.raises(ValueError):
+    not_unique = pd.DataFrame([], columns=["code", "code2", "code1", "code"])
+    with pytest.raises(
+        ValueError,
+        match="Non-unique column labels should be monotonic increasing.",
+    ):
         _select_column_names(slice("code", "code2"), not_unique)
 
 
@@ -218,11 +247,17 @@ def test_slice_dtypes(df):
     and either the start value or the stop value is not a string,
     or the step value is not an integer.
     """
-    with pytest.raises(ValueError):
+    with pytest.raises(
+        ValueError,
+        match="The start value for the slice must either be `None`.+",
+    ):
         _select_column_names(slice(1, "M_end_date_2"), df)
-    with pytest.raises(ValueError):
+    with pytest.raises(
+        ValueError,
+        match="The stop value for the slice must either be `None`.+",
+    ):
         _select_column_names(slice("id", 2), df)
-    with pytest.raises(ValueError):
+    with pytest.raises(ValueError, match="The step value for the slice.+"):
         _select_column_names(slice("id", "M_end_date_2", "3"), df)
 
 
@@ -289,52 +324,19 @@ def test_slice_dates_inexact(df_dates):
     )
 
 
-@pytest.mark.xfail(reason="level parameter removed.")
-def test_level_type(df_tuple):
-    """Raise TypeError if `level` is the wrong type."""
-    with pytest.raises(TypeError):
-        _select_column_names("A", df_tuple)
-
-
-@pytest.mark.xfail(reason="level parameter removed.")
-def test_level_nonexistent(df_tuple):
-    """
-    Raise ValueError if column is a MultiIndex
-    and level is `None`.
-    """
-    with pytest.raises(ValueError):
-        _select_column_names("A", df_tuple)
-
-
-@pytest.mark.xfail(reason="level parameter removed.")
-def test_tuple_callable(df_tuple):
-    """
-    Raise ValueError if dataframe has MultiIndex columns
-    and a callable is provided.
-    """
-    with pytest.raises(ValueError):
-        _select_column_names(lambda df: df.name.startswith("A"), df_tuple)
-
-
-@pytest.mark.xfail(reason="level parameter removed.")
-def test_tuple_regex(df_tuple):
-    """
-    Raise ValueError if dataframe has MultiIndex columns'
-    a regex is provided and level is None.
-    """
-    with pytest.raises(ValueError):
-        _select_column_names(re.compile("A"), df_tuple)
-
-
 def test_boolean_list_dtypes(df):
     """
     Raise ValueError if `columns_to_select` is a list of booleans
     and the length is unequal to the number of columns
     in the dataframe.
     """
-    with pytest.raises(ValueError):
+    with pytest.raises(
+        ValueError, match="The length of the list of booleans.+"
+    ):
         _select_column_names([True, False], df)
-    with pytest.raises(ValueError):
+    with pytest.raises(
+        ValueError, match="The length of the list of booleans.+"
+    ):
         _select_column_names(
             [True, True, True, False, False, False, True, True, True, False],
             df,
@@ -355,7 +357,9 @@ def test_callable(df_numbers):
         _select_column_names(lambda df: df + 3, df_numbers)
 
 
-@pytest.mark.xfail(reason="Indexing in Pandas is possible with a Series.")
+@pytest.mark.xfail(
+    reason="Indexing in Pandas is possible with a boolean Series."
+)
 def test_callable_returns_series(df):
     """
     Check that error is raised if `columns_to_select` is a
@@ -367,12 +371,11 @@ def test_callable_returns_series(df):
 
 def test_callable_no_match(df):
     """
-    Test if `columns_to_select` is a callable,
-    and no value is returned.
+     Raise KeyError if `columns_to_select` is a callable,
+    and no match is returned.
     """
-    assert _select_column_names(pd.api.types.is_float_dtype, df).empty
-
-    assert _select_column_names(lambda x: "Date" in x.name, df).empty
+    with pytest.raises(KeyError, match="No match was returned.+"):
+        _select_column_names(pd.api.types.is_float_dtype, df)
 
 
 def test_tuple_presence(df_tuple):