create groupby accessor and method

Author: samuel.oranyeli

.pre-commit-config.yaml

@@ -17,7 +17,7 @@ repos:
   rev: 1.5.0
   hooks:
   - id: interrogate
-    args: [-c, pyproject.toml, -vv]
+    args: [-c, pyproject.toml, -vv, --fail-under=70]
 - repo: https://github.com/terrencepreilly/darglint
   rev: v1.8.1
   hooks:

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 typing import Callable
+
+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,124 @@ 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) -> Callable:
+    """
+    Register a custom accessor on DataFrameGroupBy objects.
+
+    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: Callable) -> Callable:
+    """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 DataFrame.
+
+    Returns:
+        A Callable.
+    """
+
+    def inner(*args: tuple, **kwargs: dict):
+        class AccessorMethod(object):
+            __doc__ = method.__doc__
+
+            def __init__(self, obj):
+                self._obj = obj
+
+            @wraps(method)
+            def __call__(self, *args, **kwargs):
+                return method(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()