Python 通过方法链保留子类

Question

我定义了 pandas Dataframe的子类。该子类基本上与数据框相同，但具有完成专门任务的附加方法。

pandas 数据框最方便的属性之一是支持方法链接；也就是说，数据帧方法返回数据帧类的实例。

我希望能够使用这些方法，但是当我从子类中调用它们时，我会得到父类的实例。

import pandas as pd

class MySpecialDF(pd.DataFrame):
    def sqrt(self, colname):
        return self[colname]**2.0

df = MySpecialDF({'a':[1,2,3], 'b':[4,5,6]})
df.sqrt('a') # all good!

df = df.drop('b', axis=1) # returns a regular DF
df.sqrt('a') # AttributeError: 'DataFrame' object has no attribute 'sqrt'

我如何设置才能使这些方法返回子类的实例？

我可以手动重写单个方法，如下所示：

class MySpecialDF(pd.DataFrame):
    def sqrt(self, colname):
        return self[colname]**2.0

    def drop(self, *args, **kwargs):
        return MySpecialDF(super(MySpecialDF, self).drop(*args, **kwargs))

但 Dataframes 有很多这样的东西，我不想为每一个手动执行此操作。

我认为可能有一种方法可以应用一些装饰器来包装每个父方法，但我不确定如何做到这一点或者这是否是正确的方法。

这个问题对于子类继承返回父类实例的方法的所有情况都是普遍存在的。

有谁知道如何解决这个问题？

Answer 1

感谢@ayhan为我指明了正确的方向。我按照这个文档的评论，它展示了如何在 pandas 中具体完成子类化。修复方法是：

class MySpecialDF(pd.DataFrame):

    @property
    def _constructor(self):
        return MySpecialDF

    def sqrt(self, colname):
        return self[colname]**2.0

我不知道这是否解决了子类继承返回父类实例的方法的一般问题。但是，我不确定是否存在通用解决方案，因为返回的实例可以任意构造。