通过类似属性的重定向为 Python 类提供替代 API



这是一个通用的Python设计问题,但具体来说,我将讨论pandas.DataFrame类及其unstack方法。

我的目标是提供unstack的替代版本,但是我不想覆盖原始版本,而是提供一种通过重定向层访问我的版本的方法.x

import pandas as pd
dat = pd.DataFrame(some args)
# call my function
dat.x.unstack()
# call original function
dat.unstack()

一个设计要求是我的版本的实现应该能够调用原始unstack

为了这个问题,人们可以假设我可以向熊猫添加一些代码。数据帧,但鉴于它是一个公共包,这种添加应该是最小的。 完全非侵入性的解决方案将是一个加号。

假设你有一个名为X的类,它包含你想要在DataFrame之上实现的所有其他方法

,如下所示:
class X(object):
def __init__(self, df):
self._df = df
def unstack(self):
return 'X.untrack called!'

您可以像这样围绕DataFrame创建一个包装器:

class DataFrameWrapper(object):
def __init__(self, *args, **kwargs):
self._dataframe = pd.DataFrame(*args, **kwargs)
self._delegate = X(self._dataframe)
def __getattr__(self, attr):
if attr == 'x':
return self._delegate
return getattr(self._dataframe, attr)
dfw = DataFrameWrapper([1, 2, 3, 4, 5])
print(dfw.unstack())    # Normal unstack() output
print(dfw.x.unstack())  # X.untrack called!

还可以使用setattrDataFrame实例上的属性x设置为引用委托对象的实例。

df = pd.DataFrame([1, 2, 3, 4, 5])
setattr(df, 'x', X())
print(df.unstack())    # Normal unstack() output
print(df.x.unstack())  # X.untrack called!

编辑:将包装器中的数据帧传递给X构造函数。

这是一个受@kingkupps答案中使用getattr启发的解决方案。

import pandas as pd

class DataFrame(pd.DataFrame):
# replace pd.DataFrame.__getattr__
_super__getattr__ = pd.DataFrame.__getattr__
def _replacement__getattr__(self, name):
if name == 'x':
return DataFrame(self)
return DataFrame._super__getattr__(self, name)
pd.DataFrame.__getattr__ = _replacement__getattr__
# new api
def unstack(self, *args, **kwargs):
print('new api')
return super().unstack(*args, **kwargs)

df = pd.DataFrame(range(5))
print(df.unstack())
print(df.x.unstack())

编辑:事实证明,熊猫明确提供了一种定义此类扩展的方法,请参阅 https://pandas.pydata.org/pandas-docs/stable/development/extending.html

最新更新