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 的方法 return 另一个 DataFrame，例如上面 MWE 中的 .diff()，return 我另一个 SPDF，太棒了。但是，我还想欺骗 .resample('M').last() 或 .rolling(2).mean() 等链式方法最终生成 SPDF。到目前为止我都失败了，因为 .rolling() 等是 callable 类型，而我的包装器 to_spdf 试图从他们的输出中构造一个 SPDF 而没有 'waiting' .mean() 或表达式的任何其他最后部分。任何想法如何解决这个问题？

谢谢。

Answer 1

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

您可以执行以下操作：

class SPDF(DataFrame):

    @property
    def _constructor(self):
        return SPDF

如果您需要在 copy-constructor 方法（如 diff）期间保留自定义 attributes（不是 functions - 那些将在那里），那么您可以类似下面的内容

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'>]

Answer 2

如果你不想子class DataFrame，你可以引入另一个class像PendingSPDF并用它包裹非数据帧对象：

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)
        if isinstance(res, pd.DataFrame):
            return SPDF(res)
        else:
            return PendingSPDF(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

class PendingSPDF:
    def __init__(self, df):
        self.df = 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())

pandas：链式方法的组合，如 .resample()、.rolling() 等

pandas: Composition for chained methods like .resample(), .rolling() etc

python

chained

object-composition

pandas