Python (pandas)：使用使用 pandas API 的装饰器

Question

我对 Python 上的一般装饰器和类很陌生，但有一个问题，是否有更好的方法来装饰 Pandas 对象。举个例子，我写了以下内容来创建两个方法——lisa 和 wil：

import numpy as np
import pandas as pd

test = np.array([['john', 'meg', 2.23, 6.49],
       ['lisa', 'wil', 9.67, 8.87],
       ['lisa', 'fay', 3.41, 5.04],
       ['lisa', 'wil', 0.58, 6.12],
       ['john', 'wil', 7.31, 1.74]],
)
test = pd.DataFrame(test)
test.columns = ['name1','name2','scoreA','scoreB']

@pd.api.extensions.register_dataframe_accessor('abc')
class ABCDataFrame:

    def __init__(self, pandas_obj):
        self._obj = pandas_obj

    @property
    def lisa(self):
        return self._obj.loc[self._obj['name1'] == 'lisa']
    @property
    def wil(self):
        return self._obj.loc[self._obj['name2'] == 'wil']

示例输出如下：

test.abc.lisa.abc.wil
  name1 name2 scoreA scoreB
1  lisa   wil   9.67   8.87
3  lisa   wil   0.58   6.12

我有两个问题。

首先，在实践中，我创建的方法不止两个，并且需要在同一行中调用其中的许多方法。有没有办法test.lisa.wil返回与上面我写的相同的输出test.abc.lisa.abc.wil，因为前者将使我不必abc每次都输入？

其次，如果有关于装饰 Pandas DataFrames 的任何其他建议/资源，请告诉我。

Answer 1

您可以使用pandas-flavor库执行此操作，该库允许您DataFrame使用其他方法扩展该类。

import pandas as pd
import pandas_flavor as pf

# Create test DataFrame as before.
test = pd.DataFrame([
    ['john', 'meg', 2.23, 6.49],
    ['lisa', 'wil', 9.67, 8.87],
    ['lisa', 'fay', 3.41, 5.04],
    ['lisa', 'wil', 0.58, 6.12],
    ['john', 'wil', 7.31, 1.74]
], columns=['name1', 'name2', 'scoreA', 'scoreB'])

# Register new methods.
@pf.register_dataframe_method
def lisa(df):
    return df.loc[df['name1'] == 'lisa']

@pf.register_dataframe_method
def wil(df):
    return df.loc[df['name2'] == 'wil']

现在可以将这些视为方法，而无需中间.abc访问器。

test.lisa()                                                                                                                                                                                                                         
#   name1 name2  scoreA  scoreB
# 1  lisa   wil    9.67    8.87
# 2  lisa   fay    3.41    5.04
# 3  lisa   wil    0.58    6.12

test.lisa().wil()                                                                                                                                                                                                                   
#   name1 name2  scoreA  scoreB
# 1  lisa   wil    9.67    8.87
# 3  lisa   wil    0.58    6.12

更新

由于您有很多这样的方法，因此也可以定义一个通用的过滤方法，然后在一些循环中调用它。

def add_method(key, val, fn_name=None):  
    def fn(df):
        return df.loc[df[key] == val]

    if fn_name is None:
        fn_name = f'{key}_{val}'

    fn.__name__ = fn_name
    fn = pf.register_dataframe_method(fn)
    return fn

for name1 in ['john', 'lisa']:
    add_method('name1', name1)

for name2 in ['fay', 'meg', 'wil']:
    add_method('name2', name2)

然后这些就可以作为方法使用，就像您直接定义了方法一样。请注意，我在列名 (name1或name2)前面加上了前缀，以使其更加清晰。那是可选的。

test.name1_john()                                                                                                                                                                                                             
#   name1 name2  scoreA  scoreB
# 0  john   meg    2.23    6.49
# 4  john   wil    7.31    1.74

test.name1_lisa()                                                                                                                                                                                                                   
#   name1 name2  scoreA  scoreB
# 1  lisa   wil    9.67    8.87
# 2  lisa   fay    3.41    5.04
# 3  lisa   wil    0.58    6.12

test.name2_fay()                                                                                                                                                                                                                    
#   name1 name2  scoreA  scoreB
# 2  lisa   fay    3.41    5.04

更新 2

注册方法也可以有参数。所以另一种方法是为每列创建一个这样的方法，并将值作为参数。

@pf.register_dataframe_method
def name1(df, val):
    return df.loc[df['name1'] == val]

@pf.register_dataframe_method
def name2(df, val):
    return df.loc[df['name2'] == val]

test.name1('lisa')
#   name1 name2  scoreA  scoreB
# 1  lisa   wil    9.67    8.87
# 2  lisa   fay    3.41    5.04
# 3  lisa   wil    0.58    6.12

test.name1('lisa').name2('wil')
#   name1 name2  scoreA  scoreB
# 1  lisa   wil    9.67    8.87
# 3  lisa   wil    0.58    6.12