Python 类：如何将自定义方法添加到现有对象（熊猫数据框）

Question

我想构建一个数据查看器，利用 EXCEL 来查看 pandas DataFrame 的内容。下面是一个简单的工作示例（作为函数）。但是，我不想将其称为view(df,10)，而是将其称为df.view(10) 方法，类似于使用head 的方式，例如df.head(10)。

我是 python classes 的新手。互联网上的所有示例都是定义一个新对象并为该新对象开发类。但我想我需要为现有的 pandas DataFrame 添加一个方法。我希望将此新方法存储在我自己的私有存储库中，然后我可以导入它，例如from brb import *，并能够在任意DataFrame 上使用它。

这可能吗？

import os
import numpy as np
import pandas as pd
import xlwings as xl


df = pd.DataFrame({
    'id': (1,2,3,4,5,6,7,8,9,10)*10,
    'year': tuple(np.arange(2011,2021))*10,
    'a': np.random.choice(range(2),100),
    'b': np.random.choice(range(100),100),
})


def view(df,NObs=None):
    book = xl.Book()
    book.sheets[0].range("A1").value = df[:NObs]


#  What I have:
view(df,10)   

#  What I want:
df.view(10)

Answer 1

您可以像这样定义自己的类：

import numpy as np
import pandas as pd
import xlwings as xl


class MyDataFrame:
    def __init__(self, data):
        self.df = pd.DataFrame(data)

    def view(self, NObs=None):
        book = xl.Book()
        book.sheets[0].range("A1").value = self.df[:NObs]


if __name__ == "__main__":
    df = MyDataFrame(
        {
            "id": (1, 2, 3, 4, 5, 6, 7, 8, 9, 10) * 10,
            "year": tuple(np.arange(2011, 2021)) * 10,
            "a": np.random.choice(range(2), 100),
            "b": np.random.choice(range(100), 100),
        }
    )
    df.view(10)

然后直接运行脚本或者在别处导入MyDataFrame。

或者您甚至可以将自己的方法添加到 Pandas 对象并“本地”运行它，如下所示：

import numpy as np
import pandas as pd
import xlwings as xl
from pandas.core.base import PandasObject


def view(df, NObs=None):
    book = xl.Book()
    book.sheets[0].range("A1").value = df[:NObs]


PandasObject.view = view


if __name__ == "__main__":
    df = pd.DataFrame(
        {
            "id": (1, 2, 3, 4, 5, 6, 7, 8, 9, 10) * 10,
            "year": tuple(np.arange(2011, 2021)) * 10,
            "a": np.random.choice(range(2), 100),
            "b": np.random.choice(range(100), 100),
        }
    )
    df.view(10)