From b6c3b41ad224282dbe980751c9c574eab560ffa8 Mon Sep 17 00:00:00 2001 From: "samuel.oranyeli" Date: Sun, 26 Oct 2025 21:12:38 +1100 Subject: [PATCH 1/2] add support for SeriesGroupBy --- environment.yml | 2 +- pandas_flavor/__init__.py | 12 +++-- pandas_flavor/register.py | 95 +++++++++++++++++++++++++++++++---- tests/test_pandas_register.py | 32 ++++++++++-- 4 files changed, 123 insertions(+), 18 deletions(-) diff --git a/environment.yml b/environment.yml index 70a886b..6b47d03 100644 --- a/environment.yml +++ b/environment.yml @@ -2,7 +2,7 @@ name: pandas-flavor channels: - conda-forge dependencies: - - python=3.9 + - python>=3.9 - pandas>=0.23 - xarray - pip diff --git a/pandas_flavor/__init__.py b/pandas_flavor/__init__.py index de56082..261dc1a 100644 --- a/pandas_flavor/__init__.py +++ b/pandas_flavor/__init__.py @@ -4,8 +4,10 @@ register_dataframe_method, register_series_accessor, register_series_method, - register_groupby_accessor, - register_groupby_method, + register_dataframe_groupby_accessor, + register_dataframe_groupby_method, + register_series_groupby_accessor, + register_series_groupby_method, ) from .xarray import ( register_xarray_dataarray_method, @@ -17,8 +19,10 @@ "register_series_accessor", "register_dataframe_method", "register_dataframe_accessor", - "register_groupby_accessor", - "register_groupby_method", + "register_dataframe_groupby_accessor", + "register_dataframe_groupby_method", + "register_series_groupby_accessor", + "register_series_groupby_method", "register_xarray_dataarray_method", "register_xarray_dataset_method", ] diff --git a/pandas_flavor/register.py b/pandas_flavor/register.py index 13ff260..ca20320 100644 --- a/pandas_flavor/register.py +++ b/pandas_flavor/register.py @@ -5,7 +5,7 @@ import warnings from functools import wraps -from pandas.core.groupby.generic import DataFrameGroupBy +from pandas.core.groupby.generic import DataFrameGroupBy, SeriesGroupBy from pandas.util._exceptions import find_stack_level from pandas.api.extensions import ( register_series_accessor, @@ -285,16 +285,16 @@ def __get__(self, obj, cls): return accessor_obj -def _register_accessor(name: str, cls: DataFrameGroupBy): +def _register_accessor(name: str, cls: DataFrameGroupBy | SeriesGroupBy): """ - Register a custom accessor on a DataFrameGroupBy object. + Register a custom accessor on a pandas GroupBy 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 + cls: DataFrameGroupBy|SeriesGroupBy Returns: A class decorator. @@ -319,15 +319,19 @@ def decorator(accessor): return decorator -def register_groupby_accessor(name: str): +def register_dataframe_groupby_accessor(name: str): return _register_accessor(name, DataFrameGroupBy) -def register_groupby_method(method): +def register_series_groupby_accessor(name: str): + return _register_accessor(name, SeriesGroupBy) + + +def register_dataframe_groupby_method(method): """Register a function as a method attached to the pandas DataFrameGroupBy. Example: - >>> @register_groupby_method # doctest: +SKIP + >>> @register_dataframe_groupby_method # doctest: +SKIP >>> def print_column(grp, col): # doctest: +SKIP ... '''Print the dataframe column given''' # doctest: +SKIP ... print(grp[col]) # doctest: +SKIP @@ -347,7 +351,7 @@ 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. + decorates a function with register_dataframe_groupby_method. Args: *args: The arguments to pass to the registered method. @@ -390,7 +394,80 @@ def __call__(self, *args, **kwargs): method, method_signature, self._obj, args, kwargs ) - register_groupby_accessor(method.__name__)(AccessorMethod) + register_dataframe_groupby_accessor(method.__name__)(AccessorMethod) + return method + + return inner() + + +def register_series_groupby_method(method): + """Register a function as a method attached to the pandas SeriesGroupBy. + + Example: + >>> @register_series_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.8.0" + + Args: + method: Function to be registered as a method + on the SeriesGroupBy 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_series_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): + """SeriesGroupBy Accessor method class.""" + + __doc__ = method.__doc__ + + def __init__(self, obj): + """Initialize the accessor method class. + + Args: + obj: The pandas SeriesGroupBy 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_series_groupby_accessor(method.__name__)(AccessorMethod) return method return inner() diff --git a/tests/test_pandas_register.py b/tests/test_pandas_register.py index 8eb1a05..a2e1eca 100644 --- a/tests/test_pandas_register.py +++ b/tests/test_pandas_register.py @@ -2,7 +2,7 @@ import pandas_flavor as pf import pandas as pd -from pandas.core.groupby.generic import DataFrameGroupBy +from pandas.core.groupby.generic import DataFrameGroupBy, SeriesGroupBy def test_register_dataframe_method(): @@ -43,10 +43,10 @@ def dummy_func(s: pd.Series) -> pd.Series: ser.dummy_func() -def test_register_groupby_method(): - """Test register_groupby_method.""" +def test_register_df_groupby_method(): + """Test register_dataframe_groupby_method.""" - @pf.register_groupby_method + @pf.register_dataframe_groupby_method def dummy_func(by: DataFrameGroupBy) -> DataFrameGroupBy: """Dummy func. @@ -66,3 +66,27 @@ def dummy_func(by: DataFrameGroupBy) -> DataFrameGroupBy: ) by = df.groupby("Animal") by.dummy_func() + + +def test_register_ser_groupby_method(): + """Test register_series_groupby_method.""" + + @pf.register_series_groupby_method + def dummy_func(by: SeriesGroupBy) -> SeriesGroupBy: + """Dummy func. + + Args: + by: A SeriesGroupBy object. + + Returns: + SeriesGroupBy. + """ + return by + + df = pd.Series( + { + "Animal": ["Falcon"], + } + ) + by = df.groupby([0]) + by.dummy_func() From 2256009f2d427f2668651edf768c88fbe35f021e Mon Sep 17 00:00:00 2001 From: "samuel.oranyeli" Date: Sun, 26 Oct 2025 21:14:22 +1100 Subject: [PATCH 2/2] changelog --- CHANGELOG.md | 1 + pandas_flavor/register.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 07ef24c..671cf45 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ # Changelog ## [Unreleased] +- rename register_groupby_method to register_dataframe_groupby_method. Added register_series_groupby_method. @samukweku ## [v0.7.0] - 2025-04-11 diff --git a/pandas_flavor/register.py b/pandas_flavor/register.py index ca20320..c97466c 100644 --- a/pandas_flavor/register.py +++ b/pandas_flavor/register.py @@ -336,7 +336,7 @@ def register_dataframe_groupby_method(method): ... '''Print the dataframe column given''' # doctest: +SKIP ... print(grp[col]) # doctest: +SKIP - !!! info "New in version 0.7.0" + !!! info "New in version 0.8.0" Args: method: Function to be registered as a method