注册 Pandas“访问器”而无需重复自己

Question

我想使用 Pandas

register_series_accessor

和

register_dataframe_accessor

函数将大量函数/特性“附加”到

Series

和

DataFrame

对象。但我想以一种“不重复自己”的方式这样做。

至少在我看来，Pandas 本身实现像

pd.Series.str

这样的访问器的方式涉及编写大量样板文件......本质上是编写一个实际完成工作的“核心”函数，然后在访问器类中单独重写一个方法然后调用核心函数。我想避免这种情况。

假设我有一个对

pd.Series

对象进行操作的函数集合（实际上，我的函数比这个更有趣，而且数量更多）：

def mf_cumsum_plus_const(ser, const=0):
    """Docstring for cumsum_plus_const"""
    return ser.cumsum().add(const)


def mf_cumprod_plus_const(ser, const=0):
    """Docstring for cumprod_plus_const"""
    return ser.cumprod().add(const)

MYFUNC_LIST = [mf_cumsum_plus_const, mf_cumprod_plus_const]

我想做这样的事情：

@pd.api.extensions.register_series_accessor('mf')
class MFSeriesAccessor:
    def __init__(self, pandas_obj):
        self._obj = pandas_obj

        for func in MYFUNC_LIST:
            self.tack_on_function(func)

    def tack_on_function(self, func):
        # SOME MAGIC GOES HERE

我正在寻找一些关于在

SOME MAGIC GOES HERE

位置写什么的帮助。我认为初稿可能如下所示：

def tack_on_function(self, func):
    new_name = '_'.join(func.__name__.split('_')[1:]) # just tidying up the prefixes

    def inner(*args, **kwargs):
        return func(self._obj, *args, **kwargs)

    inner.__doc__ = func.__doc__
    inner.__name__ = new_name
    
    setattr(self, new_name, inner)

这基本上有效。但我认为这最终会产生

ser.mf.cumsum_plus_const

的“通用”函数签名。我怎样才能（以编程方式）获得该函数的签名看起来正确？

我也希望以编程方式让这些函数的

pd.DataFrame

版本正常工作。 IE。我想做一些类似的事情：

@pd.api.extensions.register_dataframe_accessor('mf')
class MFDataFrameAccessor:
    def __init__(self, pandas_obj):
        self._obj = pandas_obj

        for func in MYFUNC_LIST:
            self.tack_on_function(func)

    def tack_on_function(self, func):
        new_name = '_'.join(func.__name__.split('_')[1:]) # just tidying up the prefixes

        def inner(*args, **kwargs, axis=0):
            # note that this is doing `pd.DataFrame.apply`
            return self._obj.apply(func, axis=axis, *args, **kwargs)

        inner.__doc__ = func.__doc__
        inner.__name__ = new_name
    
        setattr(self, new_name, inner)

因此，为了让

DataFrame

案例正常工作，我需要找到一种方法来设置

inner

的签名，使其不与

func

的签名完全匹配（正如我在

Series

案例中所做的那样） ）但是，改为“匹配

func

的签名，但还包含额外的关键字参数

axis

”。

执行此操作的正确/最佳方法是什么（在

Series

和

DataFrame

情况下）？

Answer 1

以下是如何从函数列表中构造

pd.Series

和

pd.Dataframe

的访问器的示例：

from functools import partial


def mf_cumsum_plus_const(ser, const=0):
    """Docstring for cumsum_plus_const"""
    return ser.cumsum().add(const)


def mf_cumprod_plus_const(ser, const=0):
    """Docstring for cumprod_plus_const"""
    return ser.cumprod().add(const)


MYFUNC_LIST = [mf_cumsum_plus_const, mf_cumprod_plus_const]


@pd.api.extensions.register_series_accessor("mf")
class MFSeriesAccessor:
    def __init__(self, obj):
        self._obj = obj
        self._funcs = {
            f.__name__.split("_", maxsplit=1)[-1]: partial(f, self._obj)
            for f in MYFUNC_LIST
        }

    def __getattr__(self, name):
        return self._funcs[name]


@pd.api.extensions.register_dataframe_accessor("mf")
class MFDataFrameAccessor:
    def __init__(self, obj):
        self._obj = obj

        def _build_lambda(f):
            return lambda *args, **kwargs: self._obj.apply(
                f, axis=kwargs.pop("axis", 0), args=args, **kwargs
            )

        self._funcs = {
            f.__name__.split("_", maxsplit=1)[-1]: _build_lambda(f) for f in MYFUNC_LIST
        }

    def __getattr__(self, name):
        return self._funcs[name]

然后（使用示例数据框）：

print(df)

print(df.Col1.mf.cumsum_plus_const(1000))
print(df.Col1.mf.cumprod_plus_const(10))

print(df.mf.cumsum_plus_const(1000))
print(df.mf.cumprod_plus_const(10))

print(df.mf.cumsum_plus_const(1000, axis=1))
print(df.mf.cumprod_plus_const(10, axis=1))

打印：

# example df:
   Col1  Col2
0     1     2
1     3     4
2     5     6

# df.Col1.mf.cumsum_plus_const(1000)
0    1001
1    1004
2    1009
Name: Col1, dtype: int64

# df.Col1.mf.cumprod_plus_const(10)
0    11
1    13
2    25
Name: Col1, dtype: int64

# df.mf.cumsum_plus_const(1000)
   Col1  Col2
0  1001  1002
1  1004  1006
2  1009  1012

# df.mf.cumprod_plus_const(10)
   Col1  Col2
0    11    12
1    13    18
2    25    58

# df.mf.cumsum_plus_const(1000, axis=1)
   Col1  Col2
0  1001  1003
1  1003  1007
2  1005  1011

# df.mf.cumprod_plus_const(10, axis=1)
   Col1  Col2
0    11    12
1    13    22
2    15    40