Pandas DataFrameのサブクラス

DataFrameに格納されたデータに対してルーティンで施す処理が多いので、DataFrameのサブクラスを作成して、そこにメソッドをつけようとしたら苦戦したのでメモ

サブクラス作成

まずは必要なライブラリをimport

import pandas as pd
import numpy as np
import copy

クラスを作成。単純なmethodをつける。

class MyDataFrame(pd.DataFrame):
    def __init__(self,data,school):
        super( MyDataFrame,self).__init__(data)
        self.school=school
        
    def showschool(self):
        print(self.school)

インスタンスを作成し、操作してみる。

data=[['Kumiko', 1, 'euphonium'], ['Reina', 1, 'trumpet'], ['Asuka', 3, 'euphonium'], ['Yuko', 2, 'trumpet']]
mdf=MyDataFrame(data,'Kitauji')
mdf.index=['Omae', 'Kosaka', 'Tanaka', 'Yoshikawa']
mdf.columns=['name', 'grade', 'instrument']
print(mdf)
print(type(mdf))
mdf.showschool()

出力

            name  grade instrument
Omae       Kumiko      1  euphonium
Kosaka      Reina      1    trumpet
Tanaka      Asuka      3  euphonium
Yoshikawa    Yuko      2    trumpet
<class '__main__.MyDataFrame'>
Kitauji

上手くいっていそうな気がする。

ところが

DataFrame.copy()を流用すると、

mdf2=mdf.copy()
print(type(mdf2))

出力

<class 'pandas.core.frame.DataFrame'>

なんと（当たり前と言えば当たり前だが）、元のDataFrameクラスのインスタンスに戻っている。DataFrame.copy()はDataFrameクラスのインスタンスを返すようにできているからだろう。

必要に応じて自作クラスを返すように定義するしかないらしい。

例えば、

class MyDataFrame(pd.DataFrame):
    def __init__(self,data,school):
        super( MyDataFrame,self).__init__(data)
        self.school=school
        
    def showschool(self):
        print(self.school)
        
    def copy(self):
        cp=MyDataFrame(data=self, school=self.school)
        attr=copy.deepcopy(self.__dict__)
        cp.__dict__.update(attr)
        return cp
    
    def __copy__(self):
        return self.copy()
    
    def __deepcopy__(self,memo):
        return self.copy()

これで、

mdf2=mdf.copy()
print(type(mdf2))
mdf2.showschool()

出力

<class '__main__.MyDataFrame'>
Kitauji

上手くいったようだ。

locやilocも同様にしたいのだけど、もともとどういう実装になっているのかわからなくて調査中。いいやり方があったら教えてほしい。