Merge pull request #36 from pyjanitor-devs/samukweku/pandas_groupby_flavor

[ENH] Add groupby accessor and method

Author: samukweku

.github/workflows/tests.yml

@@ -18,18 +18,14 @@ jobs:
 
       # See: https://github.com/marketplace/actions/setup-miniconda
       - name: Setup miniconda
-        uses: conda-incubator/setup-miniconda@v2
+        uses: conda-incubator/setup-miniconda@v3
         with:
           auto-update-conda: true
-          miniforge-variant: Mambaforge
-          channels: conda-forge
-          activate-environment: pandas-flavor
           environment-file: environment.yml
-          use-mamba: true
+          miniforge-version: latest
 
       - name: Run unit tests
         run: |
-          conda activate pandas-flavor
           python -m pip install -e .
           pytest
 

README.md

@@ -133,6 +133,8 @@ to allow visualization of [pyjanitor](https://github.com/pyjanitor-devs/pyjanito
 - **register_dataframe_accessor**: register an accessor (and it's methods) with a pandas DataFrame.
 - **register_series_method**: register a methods directly with a pandas Series.
 - **register_series_accessor**: register an accessor (and it's methods) with a pandas Series.
+- **register_groupby_method**: register a methods directly with a pandas DataFrameGroupBy object.
+- **register_groupby_accessor**: register an accessor (and it's methods) with a pandas DataFrameGroupBy object.
 
 ## Installation
 

pandas_flavor/__init__.py

@@ -4,6 +4,8 @@
     register_dataframe_method,
     register_series_accessor,
     register_series_method,
+    register_groupby_accessor,
+    register_groupby_method,
 )
 from .xarray import (
     register_xarray_dataarray_method,
@@ -15,6 +17,8 @@
     "register_series_accessor",
     "register_dataframe_method",
     "register_dataframe_accessor",
+    "register_groupby_accessor",
+    "register_groupby_method",
     "register_xarray_dataarray_method",
     "register_xarray_dataset_method",
 ]

pandas_flavor/register.py

@@ -1,5 +1,12 @@
 """Register functions as methods of Pandas DataFrame and Series."""
+
+from __future__ import annotations
+
+import warnings
 from functools import wraps
+
+from pandas.core.groupby.generic import DataFrameGroupBy
+from pandas.util._exceptions import find_stack_level
 from pandas.api.extensions import (
     register_series_accessor,
     register_dataframe_accessor,
@@ -228,3 +235,162 @@ def __call__(self, *args, **kwargs):
         return method
 
     return inner()
+
+
+# variant of pandas' accessor
+
+# copied from pandas' accessor file - pandas/pandas/core/accessor.py
+"""
+
+accessor.py contains base classes for implementing accessor properties
+that can be mixed into or pinned onto other pandas classes.
+
+"""
+
+
+class CachedAccessor:
+    """
+    Custom property-like object.
+
+    A descriptor for caching accessors.
+
+    Parameters
+    ----------
+    name : str
+        Namespace that will be accessed under, e.g. ``df.foo``.
+    accessor : DataFrameGroupBy
+        Class with the extension methods.
+
+    Notes
+    -----
+    For accessor, The class's __init__ method assumes that one of
+    ``Series``, ``DataFrame`` or ``Index`` as the
+    single argument ``data``.
+    """
+
+    def __init__(self, name: str, accessor: DataFrameGroupBy) -> None:
+        self._name = name
+        self._accessor = accessor
+
+    def __get__(self, obj, cls):
+        if obj is None:
+            # we're accessing the attribute of the class, i.e., Dataset.geo
+            return self._accessor
+        accessor_obj = self._accessor(obj)
+        # Replace the property with the accessor object. Inspired by:
+        # https://www.pydanny.com/cached-property.html
+        # We need to use object.__setattr__ because we overwrite __setattr__ on
+        # NDFrame
+        object.__setattr__(obj, self._name, accessor_obj)
+        return accessor_obj
+
+
+def _register_accessor(name: str, cls: DataFrameGroupBy):
+    """
+    Register a custom accessor on a DataFrameGroupBy object.
+
+    Args:
+        name : str
+            Name under which the accessor should be registered.
+            A warning is issued
+            if this name conflicts with a preexisting attribute.
+        cls: DataFrameGroupBy
+
+    Returns:
+        A class decorator.
+    """
+
+    def decorator(accessor):
+        if hasattr(cls, name):
+            warnings.warn(
+                f"registration of accessor {repr(accessor)} under name "
+                f"{repr(name)} for type {repr(cls)} "
+                "is overriding a preexisting "
+                f"attribute with the same name.",
+                UserWarning,
+                stacklevel=find_stack_level(),
+            )
+        setattr(cls, name, CachedAccessor(name, accessor))
+        if not hasattr(cls, "_accessors"):
+            cls._accessors = set()
+        cls._accessors.add(name)
+        return accessor
+
+    return decorator
+
+
+def register_groupby_accessor(name: str):
+    return _register_accessor(name, DataFrameGroupBy)
+
+
+def register_groupby_method(method):
+    """Register a function as a method attached to the pandas DataFrameGroupBy.
+
+    Example:
+        >>> @register_groupby_method # doctest: +SKIP
+        >>> def print_column(grp, col): # doctest: +SKIP
+        ...    '''Print the dataframe column given''' # doctest: +SKIP
+        ...    print(grp[col]) # doctest: +SKIP
+
+    !!! info "New in version 0.7.0"
+
+    Args:
+        method: Function to be registered as a method
+            on the DataFrameGroupBy object.
+
+    Returns:
+        callable: The original method.
+    """
+    method_signature = inspect.signature(method)
+
+    def inner(*args: tuple, **kwargs: dict):
+        """Inner function to register the method.
+
+        This function is called when the user
+        decorates a function with register_groupby_method.
+
+        Args:
+            *args: The arguments to pass to the registered method.
+            **kwargs: The keyword arguments to pass to the registered method.
+
+        Returns:
+            method: The original method.
+        """
+
+        class AccessorMethod(object):
+            """DataFrameGroupBy Accessor method class."""
+
+            __doc__ = method.__doc__
+
+            def __init__(self, obj):
+                """Initialize the accessor method class.
+
+                Args:
+                    obj: The pandas DataFrameGroupBy object.
+                """
+                self._obj = obj
+
+            @wraps(method)
+            def __call__(self, *args, **kwargs):
+                """Call the accessor method.
+
+                Args:
+                    *args: The arguments to pass to the registered method.
+                    **kwargs: The keyword arguments to pass
+                        to the registered method.
+
+                Returns:
+                    object: The result of calling of the method.
+                """
+                global method_call_ctx_factory
+                if method_call_ctx_factory is None:
+                    return method(self._obj, *args, **kwargs)
+
+                return handle_pandas_extension_call(
+                    method, method_signature, self._obj, args, kwargs
+                )
+
+        register_groupby_accessor(method.__name__)(AccessorMethod)
+        return method
+
+    return inner()

tests/test_pandas_register.py

@@ -1,6 +1,8 @@
 """Tests for pandas series and dataframe method registration."""
+
 import pandas_flavor as pf
 import pandas as pd
+from pandas.core.groupby.generic import DataFrameGroupBy
 
 
 def test_register_dataframe_method():
@@ -39,3 +41,28 @@ def dummy_func(s: pd.Series) -> pd.Series:
 
     ser = pd.Series()
     ser.dummy_func()
+
+
+def test_register_groupby_method():
+    """Test register_groupby_method."""
+
+    @pf.register_groupby_method
+    def dummy_func(by: DataFrameGroupBy) -> DataFrameGroupBy:
+        """Dummy func.
+
+        Args:
+            by: A DataFrameGroupBy object.
+
+        Returns:
+            DataFrameGroupBy.
+        """
+        return by
+
+    df = pd.DataFrame(
+        {
+            "Animal": ["Falcon"],
+            "Max Speed": [380.0],
+        }
+    )
+    by = df.groupby("Animal")
+    by.dummy_func()