Repository: pyjanitor-devs/pyjanitor

Issue:
Add slice feature to `select_columns`
# Brief Description

I love the [select_columns]() method. It is powerful, especially the glob style name selection (beautiful!). I wonder if it is possible to add slice option, or if it is just unnecessary.

# Example API
Currently, you pass the selection to a list:

```python

# current implementation
df = pd.DataFrame(...).select_columns(['a', 'b',  'c', 'col_*'], invert=True)
```
I think it would be nice if we could do this:

```python
# proposed addition
df = pd.DataFrame(...).select_columns(['a': 'c', 'col_*'], invert=True)
```
Probably similar to what you do with [np.r_](https://numpy.org/doc/stable/reference/generated/numpy.r_.html)  (which is only for indices)





Discussion:

Resolving patch:
diff --git a/janitor/functions.py b/janitor/functions.py
--- a/janitor/functions.py
+++ b/janitor/functions.py
@@ -1,13 +1,11 @@
 """ General purpose data cleaning functions. """
 
-from multipledispatch import dispatch
 import collections
 import datetime as dt
 import inspect
 import re
 import unicodedata
 import warnings
-from fnmatch import translate
 from functools import partial, reduce
 from typing import (
     Any,
@@ -27,6 +25,7 @@
 import numpy as np
 import pandas as pd
 import pandas_flavor as pf
+from multipledispatch import dispatch
 from natsort import index_natsorted
 from pandas.api.types import union_categoricals
 from pandas.errors import OutOfBoundsDatetime
@@ -48,6 +47,7 @@
     _pivot_longer_pattern_match,
     _replace_empty_string_with_none,
     _replace_original_empty_string_with_none,
+    _select_columns,
     _strip_underscores,
     check,
     check_column,
@@ -793,7 +793,7 @@ def _label_encode(df, column_names):
     return df
 
 
-@dispatch(pd.DataFrame, str)
+@dispatch(pd.DataFrame, str)  # noqa: F811
 def _label_encode(df, column_names):  # noqa: F811
     le = LabelEncoder()
     check_column(df, column_names=column_names, present=True)
@@ -1141,7 +1141,7 @@ def _fill_empty(df, column_names, value=None):
     return df.fillna(value=fill_mapping)
 
 
-@dispatch(pd.DataFrame, str)
+@dispatch(pd.DataFrame, str)  # noqa: F811
 def _fill_empty(df, column_names, value=None):  # noqa: F811
     """Fill empty function for the case that column_names is a string."""
     fill_mapping = {column_names: value}
@@ -3061,7 +3061,9 @@ def currency_column_to_numeric(
 @pf.register_dataframe_method
 @deprecated_alias(search_cols="search_column_names")
 def select_columns(
-    df: pd.DataFrame, search_column_names: List[str], invert: bool = False
+    df: pd.DataFrame,
+    search_column_names: Union[str, callable, Pattern, slice, list],
+    invert: bool = False,
 ) -> pd.DataFrame:
     """Method-chainable selection of columns.
 
@@ -3069,58 +3071,160 @@ def select_columns(
 
     Optional ability to invert selection of columns available as well.
 
+    Examples
+    --------
+
+    ::
+
+        import pandas as pd
+        import numpy as np
+        import datetime
+        import re
+
+        df = pd.DataFrame(
+                {
+                    "id": [0, 1],
+                    "Name": ["ABC", "XYZ"],
+                    "code": [1, 2],
+                    "code1": [4, np.nan],
+                    "code2": ["8", 5],
+                    "type": ["S", "R"],
+                    "type1": ["E", np.nan],
+                    "type2": ["T", "U"],
+                    "code3": pd.Series(["a", "b"], dtype="category"),
+                    "type3": pd.to_datetime([np.datetime64("2018-01-01"),
+                                            datetime.datetime(2018, 1, 1)]),
+                }
+            )
+
+        df
+
+           id Name  code  code1 code2 type type1 type2 code3    type3
+        0   0  ABC     1    4.0     8    S     E     T     a 2018-01-01
+        1   1  XYZ     2    NaN     5    R   NaN     U     b 2018-01-01
+
+
+    - Select by string::
+
+        df.select_columns("id")
+           id
+       0   0
+       1   1
+
+    Select via shell-like glob strings (*) is possible::
+
+        df.select_columns("*type*")
+
+           type type1 type2      type3
+        0    S     E     T 2018-01-01
+        1    R   NaN     U 2018-01-01
+
+    - Select by slice::
+
+        df.select_columns(slice("code1", "type1"))
+
+           code1 code2 type type1
+        0    4.0     8    S     E
+        1    NaN     5    R   NaN
+
+    - Select by callable (the callable is applied to every column
+      and should return a single ``True`` or ``False`` per column)::
+
+        df.select_columns(pd.api.types.is_datetime64_dtype)
+
+               type3
+        0 2018-01-01
+        1 2018-01-01
+
+        df.select_columns(lambda x: x.name.startswith("code") or
+                                    x.name.endswith("1"))
+
+           code  code1 code2 type1 code3
+        0     1    4.0     8     E     a
+        1     2    NaN     5   NaN     b
+
+        df.select_columns(lambda x: x.isna().any())
+
+             code1 type1
+        0    4.0     E
+        1    NaN   NaN
+
+    - Select by regular expression::
+
+        df.select_columns(re.compile("\\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")])
+
+           id  code  code1 code2 code3
+        0   0     1    4.0     8     a
+        1   1     2    NaN     5     b
+
+    - Setting ``invert`` to ``True``
+      returns the complement of the columns provided::
+
+        df.select_columns(["id", "code*", slice("code", "code2")],
+                          invert = True)
+
+           Name type type1 type2      type3
+        0  ABC    S     E     T 2018-01-01
+        1  XYZ    R   NaN     U 2018-01-01
+
+    Functional usage example::
+
+       import pandas as pd
+       import janitor as jn
+
+       df = pd.DataFrame(...)
+
+       df = jn.select_columns(['a', 'b', 'col_*'],
+                              invert=True)
+
     Method-chaining example:
 
     .. code-block:: python
 
-        df = pd.DataFrame(...).select_columns(['a', 'b', 'col_*'], invert=True)
+        df = (pd.DataFrame(...)
+              .select_columns(['a', 'b', 'col_*'],
+              invert=True))
 
     :param df: A pandas DataFrame.
-    :param search_column_names: A list of column names or search strings to be
-        used to select. Valid inputs include:
-        1) an exact column name to look for
-        2) a shell-style glob string (e.g., `*_thing_*`)
+    :param search_column_names: Valid inputs include:
+
+        - an exact column name to look for
+        - a shell-style glob string (e.g., `*_thing_*`)
+        - a regular expression
+        - a callable which is applicable to each Series in the dataframe
+        - a list of all the aforementioned options.
     :param invert: Whether or not to invert the selection.
-        This will result in selection of the complement of the columns
+        This will result in the selection of the complement of the columns
         provided.
     :returns: A pandas DataFrame with the specified columns selected.
-    :raises TypeError: if input is not passed as a list.
     :raises NameError: if one or more of the specified column names or
         search strings are not found in DataFrame columns.
-    """
-    if not isinstance(search_column_names, list):
-        raise TypeError(
-            "Column name(s) or search string(s) must be passed as list"
-        )
-
-    wildcards = {col for col in search_column_names if "*" in col}
-    non_wildcards = set(search_column_names) - wildcards
-
-    if not non_wildcards.issubset(df.columns):
-        nonexistent_column_names = non_wildcards.difference(df.columns)
-        raise NameError(
-            f"{list(nonexistent_column_names)} missing from DataFrame"
-        )
 
-    missing_wildcards = []
-    full_column_list = []
+    .. # noqa: DAR402
+    """
 
-    for col in search_column_names:
-        search_string = translate(col)
-        columns = [col for col in df if re.match(search_string, col)]
-        if len(columns) == 0:
-            missing_wildcards.append(col)
-        else:
-            full_column_list.extend(columns)
+    check(
+        "search_column_names",
+        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])
 
-    if len(missing_wildcards) > 0:
-        raise NameError(
-            f"Search string(s) {missing_wildcards} not found in DataFrame"
-        )
+    full_column_list = _select_columns(search_column_names, df)
 
-    return (
-        df.drop(columns=full_column_list) if invert else df[full_column_list]
-    )
+    if invert:
+        return df.drop(columns=full_column_list)
+    return df.loc[:, full_column_list]
 
 
 @pf.register_dataframe_method
diff --git a/janitor/utils.py b/janitor/utils.py
--- a/janitor/utils.py
+++ b/janitor/utils.py
@@ -1,10 +1,13 @@
 """Miscellaneous internal PyJanitor helper functions."""
 
+import fnmatch
 import functools
 import os
+import re
 import sys
 import warnings
 from collections import namedtuple
+from collections.abc import Callable as dispatch_callable
 from itertools import chain, product
 from typing import (
     Callable,
@@ -28,6 +31,7 @@
 def check(varname: str, value, expected_types: list):
     """
     One-liner syntactic sugar for checking types.
+    It can also check callables.
 
     Should be used like this::
 
@@ -38,10 +42,13 @@ def check(varname: str, value, expected_types: list):
     :param expected_types: The types we expect the item to be.
     :raises TypeError: if data is not the expected type.
     """
-    is_expected_type = False
+    is_expected_type: bool = False
     for t in expected_types:
-        if isinstance(value, t):
-            is_expected_type = True
+        if t is callable:
+            is_expected_type = t(value)
+        else:
+            is_expected_type = isinstance(value, t)
+        if is_expected_type:
             break
 
     if not is_expected_type:
@@ -1680,3 +1687,181 @@ def _computations_as_categorical(df: pd.DataFrame, **kwargs) -> pd.DataFrame:
     df = df.astype(categories_dtypes)
 
     return df
+
+
+@functools.singledispatch
+def _select_columns(columns_to_select: str, df):
+    """
+    Base function for column selection.
+    Applies only to strings.
+    It is also applicable to shell-like glob strings,
+    specifically, the `*`.
+    A list/pd.Index of column names is returned.
+    """
+    filtered_columns = None
+    if "*" in columns_to_select:  # shell-style glob string (e.g., `*_thing_*`)
+        filtered_columns = fnmatch.filter(df.columns, columns_to_select)
+    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}'")
+    return filtered_columns
+
+
+@_select_columns.register(slice)  # noqa: F811
+def _column_sel_dispatch(columns_to_select, df):  # noqa: F811
+    """
+    Base function for column selection.
+    Applies only to slices.
+    The start slice value must be a string or None;
+    same goes for the stop slice value.
+    The step slice value should be an integer or None.
+    A list/pd.Index of column names is returned.
+    """
+    filtered_columns = None
+    start_check = None
+    stop_check = None
+    step_check = None
+    # df.columns should be monotonic/unique,
+    # but we wont check for that
+    # onus is on the user to ensure that
+    start, stop, step = (
+        columns_to_select.start,
+        columns_to_select.stop,
+        columns_to_select.step,
+    )
+    start_check = any((start is None, isinstance(start, str)))
+    stop_check = any((stop is None, isinstance(stop, str)))
+    step_check = any((step is None, isinstance(step, int)))
+    if not start_check:
+        raise ValueError(
+            """
+            The start value for the slice
+            must either be a string or `None`.
+            """
+        )
+    if not stop_check:
+        raise ValueError(
+            """
+            The stop value for the slice
+            must either be a string or `None`.
+            """
+        )
+    if not step_check:
+        raise ValueError(
+            """
+            The step value for the slice
+            must either be an integer or `None`.
+            """
+        )
+    start_check = any((start is None, start in df.columns))
+    stop_check = any((stop is None, stop in df.columns))
+    if not start_check:
+        raise ValueError(
+            """
+            The start value for the slice must either be `None`
+            or exist in the dataframe's columns.
+            """
+        )
+    if not stop_check:
+        raise ValueError(
+            """
+            The stop value for the slice must either be `None`
+            or exist in the dataframe's columns.
+            """
+        )
+    filtered_columns = df.columns.slice_locs(start=start, end=stop)
+    # slice_locs fails when step has a value
+    # so this extra step is necessary to get the correct output
+    start, stop = filtered_columns
+    filtered_columns = df.columns[slice(start, stop, step)]
+    return filtered_columns
+
+
+@_select_columns.register(dispatch_callable)  # noqa: F811
+def _column_sel_dispatch(columns_to_select, df):  # noqa: F811
+    """
+    Base function for column selection.
+    Applies only to callables.
+    The callable is applied to every column in the dataframe.
+    Either True or False is expected per column.
+    A list/pd.Index of column names is returned.
+    """
+    # the function will be applied per series.
+    # this allows filtration based on the contents of the series
+    # or based on the name of the series,
+    # which happens to be a column name as well.
+    # whatever the case may be,
+    # the returned values should be a sequence of booleans,
+    # with at least one True.
+    filtered_columns = None
+    filtered_columns = [columns_to_select(column) for _, column in df.items()]
+
+    # returns numpy bool,
+    # which does not work the same way as python's bool
+    # as such, cant use isinstance.
+    # pandas type check function helps out
+    checks = (pd.api.types.is_bool(column) for column in filtered_columns)
+
+    if not all(checks):
+        raise ValueError(
+            "The callable provided must return a sequence of booleans."
+        )
+
+    # cant use `is` here, since it may be a numpy bool
+    checks = any((column == True for column in filtered_columns))  # noqa: E712
+
+    if not checks:
+        raise ValueError("No results were returned for the callable provided.")
+
+    filtered_columns = df.columns[filtered_columns]
+    return filtered_columns
+
+
+# hack to get it to recognize typing.Pattern
+# functools.singledispatch does not natively
+# recognize types from the typing module
+# ``type(re.compile(r"\d+"))`` returns re.Pattern
+# which is a type and functools.singledispatch
+# accepts it without drama;
+# however, the same type from typing.Pattern
+# is not accepted.
+@_select_columns.register(type(re.compile(r"\d+")))  # noqa: F811
+def _column_sel_dispatch(columns_to_select, df):  # noqa: F811
+    """
+    Base function for column selection.
+    Applies only to regular expressions.
+    `re.compile` is required for the regular expression.
+    A list/pd.Index of column names is returned.
+    """
+    filtered_columns = None
+    filtered_columns = [
+        column_name
+        for column_name in df
+        if re.search(columns_to_select, column_name)
+    ]
+
+    if not filtered_columns:
+        raise ValueError("No column name matched the regular expression.")
+
+    return filtered_columns
+
+
+@_select_columns.register(list)  # noqa: F811
+def _column_sel_dispatch(columns_to_select, df):  # noqa: F811
+    """
+    Base function for column selection.
+    Applies only to list type.
+    It can take any of slice, str, callable, re.Pattern types,
+    or a combination of these types.
+    A list/pd.Index of column names is returned.
+    """
+    filtered_columns = []
+
+    for entry in columns_to_select:
+        outcome = [
+            c for c in _select_columns(entry, df) if c not in filtered_columns
+        ]
+        filtered_columns.extend(outcome)
+
+    return 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
@@ -50,9 +50,54 @@ def test_select_columns_missing_columns(dataframe, columns):
 
 @pytest.mark.functions
 @pytest.mark.parametrize(
-    "columns", ["a", ("a", "Bell__Chart"), {"a", "Bell__Chart"}]
+    "columns",
+    [
+        pytest.param(
+            "a",
+            marks=pytest.mark.xfail(
+                reason="select_columns now accepts other data types"
+            ),
+        ),
+        ("a", "Bell__Chart"),
+        {"a", "Bell__Chart"},
+    ],
 )
 def test_select_columns_input(dataframe, columns):
     """Check that passing an iterable that is not a list raises TypeError."""
     with pytest.raises(TypeError):
         dataframe.select_columns(search_column_names=columns)
+
+
+@pytest.mark.functions
+@pytest.mark.parametrize(
+    "invert,expected",
+    [
+        (False, ["Bell__Chart", "a", "decorated-elephant"]),
+        (True, ["animals@#$%^", "cities"]),
+    ],
+)
+def test_select_unique_columns(dataframe, invert, expected):
+    """Test that only unique columns are returned."""
+    columns = ["Bell__*", slice("a", "decorated-elephant")]
+    df = dataframe.select_columns(search_column_names=columns, invert=invert)
+
+    assert_frame_equal(df, dataframe[expected])
+
+
+@pytest.mark.functions
+@pytest.mark.parametrize(
+    "invert,expected",
+    [
+        (False, ["Bell__Chart", "decorated-elephant"]),
+        (True, ["a", "animals@#$%^", "cities"]),
+    ],
+)
+def test_select_callable_columns(dataframe, invert, expected):
+    """Test that columns are returned when a callable is passed."""
+
+    def columns(x):
+        return "-" in x.name or "_" in x.name
+
+    df = dataframe.select_columns(search_column_names=columns, invert=invert)
+
+    assert_frame_equal(df, dataframe[expected])
diff --git a/tests/utils/test__select_column.py b/tests/utils/test__select_column.py
new file mode 100644
--- /dev/null
+++ b/tests/utils/test__select_column.py
@@ -0,0 +1,265 @@
+import datetime
+import re
+
+import numpy as np
+import pandas as pd
+import pytest
+from pandas.testing import assert_index_equal
+
+from janitor import patterns
+from janitor.utils import _select_columns
+
+
+@pytest.fixture
+def df():
+    return pd.DataFrame(
+        {
+            "id": [1, 2, 3],
+            "M_start_date_1": [201709, 201709, 201709],
+            "M_end_date_1": [201905, 201905, 201905],
+            "M_start_date_2": [202004, 202004, 202004],
+            "M_end_date_2": [202005, 202005, 202005],
+            "F_start_date_1": [201803, 201803, 201803],
+            "F_end_date_1": [201904, 201904, 201904],
+            "F_start_date_2": [201912, 201912, 201912],
+            "F_end_date_2": [202007, 202007, 202007],
+        }
+    )
+
+
+@pytest.fixture
+def df1():
+    return pd.DataFrame(
+        {
+            "id": [0, 1],
+            "Name": ["ABC", "XYZ"],
+            "code": [1, 2],
+            "code1": [4, np.nan],
+            "code2": ["8", 5],
+            "type": ["S", "R"],
+            "type1": ["E", np.nan],
+            "type2": ["T", "U"],
+            "code3": pd.Series(["a", "b"], dtype="category"),
+            "type3": pd.to_datetime(
+                [np.datetime64("2018-01-01"), datetime.datetime(2018, 1, 1)]
+            ),
+        }
+    )
+
+
+def test_type(df):
+    """Raise TypeError if `columns_to_select` is the wrong type."""
+    with pytest.raises(TypeError):
+        _select_columns(2.5, df)
+    with pytest.raises(TypeError):
+        _select_columns(1, df)
+    with pytest.raises(TypeError):
+        _select_columns([3, "id"], df)
+
+
+def test_strings_do_not_exist(df):
+    """
+    Raise NameError if `columns_to_select` is a string
+    and does not exist in the dataframe's columns.
+    """
+    with pytest.raises(NameError):
+        _select_columns("word", df)
+    with pytest.raises(NameError):
+        _select_columns("*starter", df)
+
+
+def test_slice_dtypes(df):
+    """
+    Raise ValueError if `columns_to_select` is a slice instance
+    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):
+        _select_columns(slice(1, "M_end_date_2"), df)
+    with pytest.raises(ValueError):
+        _select_columns(slice("id", 2), df)
+    with pytest.raises(ValueError):
+        _select_columns(slice("id", "M_end_date_2", "3"), df)
+
+
+def test_slice_presence(df):
+    """
+    Raise ValueError if `columns_to_select` is a slice instance
+    and either the start value or the end value is not present
+    in the dataframe.
+    """
+    with pytest.raises(ValueError):
+        _select_columns(slice("Id", "M_start_date_1"), df)
+    with pytest.raises(ValueError):
+        _select_columns(slice("id", "M_end_date"), df)
+
+
+def test_callable(df):
+    """
+    Check that error is raised if `columns_to_select` is a
+    callable, and at lease one Series has a wrong data type
+    that makes the callable unapplicable.
+    """
+    with pytest.raises(TypeError):
+        _select_columns(object, df)
+
+
+def test_callable_returns_Series(df):
+    """
+    Check that error is raised if `columns_to_select` is a
+    callable, and returns a Series.
+    """
+    with pytest.raises(ValueError):
+        _select_columns(lambda x: x + 1, df)
+
+
+def test_callable_no_match(df):
+    """
+    Raise ValueError if `columns_to_select` is a callable, and
+    no boolean results are returned, when the callable is
+    applied to each series in the dataframe.
+    """
+    with pytest.raises(ValueError):
+        _select_columns(pd.api.types.is_float_dtype, df)
+
+    with pytest.raises(ValueError):
+        _select_columns(lambda x: "Date" in x.name, df)
+
+
+def test_regex_presence(df):
+    """
+    Raise ValueError if `columns_to_select` is a regex
+    and none of the column names match.
+    """
+    with pytest.raises(ValueError):
+        _select_columns(re.compile(r"^\d+"), df)
+
+
+def test_strings(df1):
+    """Test _select_columns function on strings."""
+    assert _select_columns("id", df1) == ["id"]
+    assert _select_columns("*type*", df1) == [
+        "type",
+        "type1",
+        "type2",
+        "type3",
+    ]
+
+
+def test_slice(df1):
+    """Test _select_columns function on slices."""
+    assert_index_equal(
+        _select_columns(slice("code", "code2"), df1),
+        df1.loc[:, slice("code", "code2")].columns,
+    )
+    assert_index_equal(
+        _select_columns(slice("code2", None), df1),
+        df1.loc[:, slice("code2", None)].columns,
+    )
+    assert_index_equal(
+        _select_columns(slice(None, "code2"), df1),
+        df1.loc[:, slice(None, "code2")].columns,
+    )
+    assert_index_equal(_select_columns(slice(None, None), df1), df1.columns)
+    assert_index_equal(
+        _select_columns(slice(None, None, 2), df1),
+        df1.loc[:, slice(None, None, 2)].columns,
+    )
+
+
+def test_callable_data_type(df1):
+    """
+    Test _select_columns function on callables,
+    specifically for data type checks.
+    """
+    assert_index_equal(
+        _select_columns(pd.api.types.is_integer_dtype, df1),
+        df1.select_dtypes(int).columns,
+    )
+    assert_index_equal(
+        _select_columns(pd.api.types.is_float_dtype, df1),
+        df1.select_dtypes(float).columns,
+    )
+    assert_index_equal(
+        _select_columns(pd.api.types.is_numeric_dtype, df1),
+        df1.select_dtypes("number").columns,
+    )
+    assert_index_equal(
+        _select_columns(pd.api.types.is_categorical_dtype, df1),
+        df1.select_dtypes("category").columns,
+    )
+    assert_index_equal(
+        _select_columns(pd.api.types.is_datetime64_dtype, df1),
+        df1.select_dtypes(np.datetime64).columns,
+    )
+    assert_index_equal(
+        _select_columns(pd.api.types.is_object_dtype, df1),
+        df1.select_dtypes("object").columns,
+    )
+
+
+def test_callable_string_methods(df1):
+    """
+    Test _select_columns function on callables,
+    specifically for column name checks.
+    """
+    assert_index_equal(
+        _select_columns(lambda x: x.name.startswith("type"), df1),
+        df1.filter(like="type").columns,
+    )
+    assert_index_equal(
+        _select_columns(lambda x: x.name.endswith(("1", "2", "3")), df1),
+        df1.filter(regex=r"\d$").columns,
+    )
+    assert_index_equal(
+        _select_columns(lambda x: "d" in x.name, df1),
+        df1.filter(regex="d").columns,
+    )
+    assert_index_equal(
+        _select_columns(
+            lambda x: x.name.startswith("code") and x.name.endswith("1"), df1
+        ),
+        df1.filter(regex=r"code.*1$").columns,
+    )
+    assert_index_equal(
+        _select_columns(
+            lambda x: x.name.startswith("code") or x.name.endswith("1"), df1
+        ),
+        df1.filter(regex=r"^code.*|.*1$").columns,
+    )
+
+
+def test_callable_computations(df1):
+    """
+    Test _select_columns function on callables,
+    specifically for computations.
+    """
+    assert_index_equal(
+        _select_columns(lambda x: x.isna().any(), df1),
+        df1.columns[df1.isna().any().array],
+    )
+
+
+def test_regex(df1):
+    """Test _select_columns function on regular expressions."""
+    assert _select_columns(re.compile(r"\d$"), df1) == list(
+        df1.filter(regex=r"\d$").columns
+    )
+    assert _select_columns(patterns(r"\d$"), df1) == list(
+        df1.filter(regex=r"\d$").columns
+    )
+
+
+def test_list_various(df1):
+    """Test _select_columns function on list type."""
+
+    assert _select_columns(["id", "Name"], df1) == ["id", "Name"]
+    assert _select_columns(["id", "code*"], df1) == list(
+        df1.filter(regex="^id|^code").columns
+    )
+    assert_index_equal(
+        pd.Index(
+            _select_columns(["id", "code*", slice("code", "code2")], df1)
+        ),
+        df1.filter(regex="^(id|code)").columns,
+    )