Repository: pyjanitor-devs/pyjanitor

Issue:
Extend `select_columns` to `DataFrame.GroupBy`
More like a helper function to conveniently select columns on a groupby; can be handy for lazy groupby.  

See this question on [SO](https://stackoverflow.com/q/63200530/7175713), it could be solved easier maybe? :

```py
dat = df.groupby("cyl")
aggs = [dat.select_columns("*p").mean(), dat.select_columns("*t").sum()]
pd.concat(aggs, axis=1)
```


Discussion:

Resolving patch:
diff --git a/janitor/functions/__init__.py b/janitor/functions/__init__.py
--- a/janitor/functions/__init__.py
+++ b/janitor/functions/__init__.py
@@ -80,4 +80,5 @@
     unionize_dataframe_categories,
     DropLabel,
     get_index_labels,
+    get_columns,
 )
diff --git a/janitor/functions/utils.py b/janitor/functions/utils.py
--- a/janitor/functions/utils.py
+++ b/janitor/functions/utils.py
@@ -17,6 +17,7 @@
     Any,
 )
 from pandas.core.dtypes.generic import ABCPandasArray, ABCExtensionArray
+from pandas.core.groupby.generic import DataFrameGroupBy, SeriesGroupBy
 from pandas.core.common import is_bool_indexer
 from dataclasses import dataclass
 from enum import Enum
@@ -28,7 +29,6 @@
     is_list_like,
     is_datetime64_dtype,
     is_string_dtype,
-    is_categorical_dtype,
     is_extension_array_dtype,
     is_bool_dtype,
 )
@@ -252,7 +252,7 @@ def _is_str_or_cat(index):
     Check if the column/index is a string,
     or categorical with strings.
     """
-    if is_categorical_dtype(index):
+    if isinstance(index.dtype, pd.CategoricalDtype):
         return is_string_dtype(index.categories)
     return is_string_dtype(index)
 
@@ -622,6 +622,27 @@ def get_index_labels(
     return index[_select_index(arg, df, axis)]
 
 
+def get_columns(group: Union[DataFrameGroupBy, SeriesGroupBy], label):
+    """
+    Helper function for selecting columns on a grouped object,
+    using the
+    [`select_columns`][janitor.functions.select.select_columns] syntax.
+
+    !!! info "New in version 0.25.0"
+
+    Args:
+        group: A Pandas GroupBy object.
+        label: column(s) to select.
+
+    Returns:
+        A pandas groupby object.
+    """
+    check("groupby object", group, [DataFrameGroupBy, SeriesGroupBy])
+    label = get_index_labels(label, group.obj, axis="columns")
+    label = label if is_scalar(label) else list(label)
+    return group[label]
+
+
 def _select(
     df: pd.DataFrame,
     args: tuple,


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
@@ -6,7 +6,7 @@
 from pandas.testing import assert_frame_equal
 from itertools import product
 
-from janitor.functions.utils import patterns, DropLabel
+from janitor.functions.utils import patterns, DropLabel, get_columns
 from pandas.api.types import is_numeric_dtype
 
 
@@ -447,3 +447,10 @@ def test_regex_multi(multiindex):
     expected = multiindex.select_columns(re.compile("b.r"))
     actual = multiindex.loc(axis=1)[["bar"]]
     assert_frame_equal(expected, actual)
+
+
+def test_select_groupby(dataframe):
+    """Test output on a grouped object"""
+    expected = dataframe.select_dtypes("number").groupby(dataframe["a"]).sum()
+    actual = dataframe.groupby("a").pipe(get_columns, is_numeric_dtype).sum()
+    assert_frame_equal(expected, actual)