2021-10-22

Python で関数をラップする

Python で関数をラップする方法を調べた。

関数の実行時と終了時を知らせる wrapper

関数の実行時にstart--、終了時にend----を出力する wrapper。

def dec(func):
    def callable(*args, **kwargs):
        print("start--")
        func(*args, **kwargs)
        print("end----")

    return callable

さらに、与えた文字列をprintするだけの関数を用意しておく。

def show(text):
		print(text)

動かしてみる。

dec(show)("hoge")

# 出力
# start--
# hoge
# end----

デコレータ式

関数の宣言直前に@<なんとか>と書くと関数をラップできる。


関数定義は一つ以上の デコレータ 式でラップできます。デコレータ式は関数を定義するとき、関数定義の入っているスコープで評価されます。その結果は、関数オブジェクトを唯一の引数にとる呼び出し可能オブジェクトでなければなりません。関数オブジェクトの代わりに、返された値が関数名に束縛されます。複数のデコレータはネストして適用されます。例えば、以下のようなコード:

@f1(arg)
@f2
def func(): pass

は、だいたい次と等価です

def func(): pass
func = f1(arg)(f2(func))

ただし、前者のコードでは元々の関数を func という名前へ一時的に束縛することはない、というところを除きます。

参考:https://docs.python.org/ja/3/reference/compound_stmts.html#function-definitions


@decshow関数につけることで、show関数を呼び出した時にstartendが表示されるようになる。

@dec
def show(text: str):
    print(text)

show("hoge")

# 出力
# start--
# hoge
# end----

表示する文字列を変える

今までの例ではstartendという文字列を表示していた。次は、その文字列を任意の文字列に変える wrapper を書く。

def dec(dec_str):
    def _dec(func):
        def callable(*args, **kwargs):
            print(dec_str)
            func(*args, **kwargs)
            print(dec_str)

        return callable

    return _dec

こんな感じに、表示する文字列を与えると wrapper となる関数を返すような関数を書けば良い。

@dec("----------")
def show(text: str):
    print(text)


@dec("**********")
def show2(text: str):
    print(text)


show("hoge")
# 出力
# ----------
# hoge
# ----------


show2("hoge")
# 出力
# **********
# hoge
# **********

継承したときの挙動

クラスを継承させるときに、親が持つメソッドがラップされている場合どうなるか。

def dec(func):
    def callable(*args, **kwargs):
        print("start--")
        func(*args, **kwargs)
        print("end----")

    return callable


class C1:
    @dec
    def show(self, text):
        print("C1", text)


class C2(C1):
    pass


C1().show("hoge")
# 出力
# start--
# C1 hoge
# end----

C2().show("hoge")
# 出力
# start--
# C1 hoge
# end----

この場合だと単にラップされたshowを呼ぶだけなので、C1C2どちらからの呼び出しでも同じ文字列が表示される。


C2クラスでshowメソッドを上書きした場合、startendは表示されない。

class C2(C1):
    def show(self, text):
        print("C2", text)
        
C2().show("hoge")
# 出力
# C2 hoge

C1を親に持つ子クラスのshowメソッドはdec関数でラップされていて欲しい、という場合は__new__が呼び出されたときにメソッドをラップされたもので置き換えれば良い。

参考:https://docs.python.org/3/reference/datamodel.html#object.__new__

class C1:
    def __new__(cls):
        cls.show = dec(cls.show)
        return super().__new__(cls)

    def show(self, text):
        print("C1", text)


class C2(C1):
    def show(self, text):
        print("C2", text)


C1().show("hoge")
# 出力
# start--
# C1 hoge
# end----

C2().show("hoge")
# 出力
# start--
# C2 hoge
# end----

クラスデコレータ

あるクラスの特定のメソッドをラップしたい、というとき、以下のようにも書ける。

def dec(cls):
    show = getattr(cls, "show")

    def new_show(*args, **kwargs):
        print("start--")
        show(*args, **kwargs)
        print("end----")

    setattr(cls, "show", new_show)
    return cls


@dec
class C1:
    def show(self, text):
        print(text)


C1().show("hoge")

わからないこと

@classmethodしか持たないクラスは__new__が呼ばれないので、__new__時にラップできない。