__init__()の引数を書き換えずに処理を追加する

Python Tips.

忘れないようにメモ。

実現したいこと

サードパーティ―ライブラリ(3rd-lib)の上に自身のライブラリを作りたい時、3rd-libのクラスを継承してサブクラスを作りたい時がある。

このとき、初期化処理を上書きしようと、

import pandas as pd

class MyDataFrame(pd.DataFrame):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        do_something()

としてしまうと、スーパークラスの初期化引数の情報が失われてしまう。 この場合で言えば、data, columnsといったpd.DataFrameの初期化引数が失われてコード補完が効かなくなって効率が落ちる。

最も単純な実装はスーパークラスと同じ引数をサブクラスの__init__()にも書くことだが、3rd-libのスーパークラスが変更されるとサブクラスも変更しなければならないためスマートな方法ではない。

そこで引数の情報を残したまま、3rd-libの仕様変更に依存しない実装方法を提示する。

実装方法

__new__()を利用する。

import pandas as pd

class MyDataFrame(pd.DataFrame):
    def __new__(cls, *args, **kwargs) -> MyDataFrame:
        self = super(MyDataFrame, cls).__new__(cls)
        self.__init__(*args, **kwargs)
        do_something()
        return self

注意点としては、__new__()の後に__init__()が再び呼ばれるため、副作用のある処理を__init__()に書いているとバグが発生する可能性があることくらいか。

挙動テスト

test.py

class Super:
    def __init__(self):
        print('super __init__')


class Sub(Super):
    def __new__(cls, *args, **kwargs):
        self = super(Sub, cls).__new__(cls)
        self.__init__(*args, **kwargs)
        print('sub __new__')
        return self

s = Sub()

実行結果

$ python test.py
super __init__
sub __new__
super __init__

改良版

super __init__を一回だけ呼ぶ。

class Super:
    def __init__(self):
        print('super __init__')


class Sub(Super):
    def __new__(cls, *args, **kwargs):
        self = super(Sub, cls).__new__(cls)
        self.__init__(*args, **kwargs)
        print('sub __new__')
        def f(*args, **kwargs) -> None:
            pass
        cls.__init__ = f
        return self

s = Sub()

実行結果

$ python test.py
super __init__
sub __new__