如何在带有修改参数的python中将外部函数包含到类中？

Question

我试图在类中包含外部定义的函数。 我已经制作了很多熊猫操作函数并且它们运行良好，但是当我尝试将它们添加到 Class 中时，我遇到了问题。这是我的大纲：

MWE

import numpy as np
import pandas as pd

df = pd.DataFrame({'A': [1,2,3],'B':[10,20,30],
                   'C':[100,200,300],'D':[1000,2000,3000]})

@pd.api.extensions.register_dataframe_accessor("my")
class MyAccessor:
    def __init__(self, pandas_obj):
        self._obj = pandas_obj

    def add_1_2(self,col1,col2):
        df = self._obj
        return df[col1] + df[col2]
    
df.my.add_1_2('A','B')

这可行（但我已经定义了没有 self._obj 的函数）

# this works
def add_1_3(self,col1,col2,col3):
    df = self._obj
    return df[col1] + df[col3]

MyAccessor.add_1_3 = add_1_3

df.my.add_1_3('A','B','C')

我的尝试

def temp_fn(df,col1,col2,mydict):
    print(mydict)
    return df[col1]+df[col2]

col1,col2,mydict = 'A','B',{'lang':'python'}

temp_fn(df,col1,col2,mydict) # calling directly using function works good

# now I want to include this function inside the Class
def make_class_fn(fn, *args, **kwargs):
    return fn(args[0]._obj, *args[1:], **kwargs)

MyAccessor.temp_fn = make_class_fn(temp_fn,df,col1,col2,mydict)

AttributeError: 'DataFrame' object has no attribute '_obj' 

注意

如果我只有一个函数，我可以在 Class 中复制相同的函数并在那里编辑它，但是我有大量的函数，这不是一个好主意。

资源

• Define a method outside of class definition?
• Python function assignment outside class definition causes argument exception

Answer 1

您的转换函数只需要创建一个以self._obj 作为第一个参数调用原始函数的新函数。

from functools import wraps


def temp_fn(df,col1,col2,mydict):
    print(mydict)
    return df[col1]+df[col2]


def make_method(f):
    @wraps(f)
    def _(self, col1, col2, mydict):
        return f(self._obj, col1, col2, mydict)
    return _

MyAccessor.temp_fn = make_method(temp_fn)

更一般地说，make_method 可以处理任意参数，只要 first 参数具有要传递的 _obj 属性。

def make_method(f):
    @wraps(f)
    def _(self, *args, **kwargs):
        return f(self._obj, *args, **kwargs)
    return _

您也可以让make_method 为您执行任务。

def make_method(cls, f):
    @wraps(f)
    def _(self, *args, **kwargs):
        return f(self._obj, *args, **kwargs)
    setattr(cls, f.__name__, _)

make_method(MyAccessor, temp_fn)

甚至可以将其设为装饰器（尽管如果您有要装饰的定义，您可以直接定义方法，同样简单）。

def make_method(cls):
    def decorator(f):
        @wraps(f)
        def _(self, *args, **kwargs):
            return f(self._obj, *args, **kwargs)
        setattr(cls, f.__name__, _)
        return _
    return decorator


@make_method(MyAccessor)
def temp_fn(df, col1, col2, mydict):
    return df[col1] + df[col2]