如何为pandas GroupBy对象实现自定义API访问器

解决方法 #

你之前仅为DataFrame注册了geo访问器，DataFrameGroupBy是独立的类，自然识别不到对应属性，pandas支持直接为分组对象注册自定义访问器，适配后即可实现你要的调用语法。

修改后完整可运行代码如下：

import pandas as pd
import numpy as np
from pandas.core.groupby.generic import DataFrameGroupBy

# 原有DataFrame访问器保持不变
@pd.api.extensions.register_dataframe_accessor("geo")
class GeoAccessor:
    def __init__(self, pandas_obj):
        self._obj = pandas_obj

    @property
    def center(self):
        lat = self._obj.latitude
        lon = self._obj.longitude
        return (float(lon.mean()), float(lat.mean()))

# 新增DataFrameGroupBy的geo访问器注册
@pd.api.extensions.register_accessor("geo", DataFrameGroupBy)
class GeoGroupByAccessor:
    def __init__(self, groupby_obj):
        self._grp = groupby_obj

    @property
    def center(self):
        # 复用已有逻辑，对每个分组调用DataFrame访问器的center属性
        return self._grp.apply(lambda x: x.geo.center)

if __name__ == "__main__":
    ds = pd.DataFrame({"longitude": np.linspace(0, 10),
                       "latitude": np.linspace(0, 20)})
    ds['grp'] = ds['longitude'].astype(int)
    # 直接调用你需要的语法即可得到结果
    print(ds.groupby('grp').geo.center)

运行后会输出和你之前循环打印结果一致的Series，每个索引对应分组的grp值，值为对应分组的中心坐标。如果需要适配SeriesGroupBy对象，用同样的方法将访问器注册到SeriesGroupBy类即可。

内容的提问来源于stack exchange，提问作者divingTobi