pandas:链接方法的组合,如.resample(),.rolling()等

Question

我想构建一个扩展pandas.DataFrame- 让我们称之为SPDF- 它可以做到超出简单的DataFrame能力:

import pandas as pd
import numpy as np


def to_spdf(func):
    """Transform generic output of `func` to SPDF.

    Returns
    -------
    wrapper : callable
    """
    def wrapper(*args, **kwargs):
        res = func(*args, **kwargs)
        return SPDF(res)

    return wrapper


class SPDF:
    """Special-purpose dataframe.

    Parameters
    ----------
    df : pandas.DataFrame

    """

    def __init__(self, df):
        self.df = df

    def __repr__(self):
        return repr(self.df)

    def __getattr__(self, item):
        res = getattr(self.df, item)

        if callable(res):
            res = to_spdf(res)

        return res


if __name__ == "__main__":

    # construct a generic SPDF
    df = pd.DataFrame(np.eye(4))
    an_spdf = SPDF(df)

    # call .diff() to obtain another SPDF
    print(an_spdf.diff())

现在,DataFrame返回另一个的方法DataFrame,例如.diff()在上面的MWE中,返回另一个SPDF,这很好.但是,我还想欺骗链式方法,例如.resample('M').last()或最终.rolling(2).mean()产生一个SPDF.到目前为止我失败了因为.rolling()类似的类型callable,我的包装器to_spdf试图SPDF从它们的输出构造一个没有"等待" .mean()或表达式的任何其他最后部分.有任何想法如何解决这个问题？

谢谢.

Answer 1

你应该正确地子类化dataframe。为了使copy-constructor方法正常工作，pandas 描述您必须设置属性_constructor（以及其他信息）。

您可以执行以下操作：

class SPDF(DataFrame):

    @property
    def _constructor(self):
        return SPDF

如果您需要在方法期间保留自定义attributes（不是functions- 那些将会在那里），那么您可以执行以下操作copy-constructordiff

class SPDF(DataFrame):
    _metadata = ['prop']
    prop = 1

    @property
    def _constructor(self):
        return SPDF

请注意，输出是所需的：

df = SPDF(np.eye(4))
print(type(df))
[<class '__main__.SPDF'>]
new = df.diff()
print(type(new))
[<class '__main__.SPDF'>]