tr_ikym_blog

某Webサービスに従事する者のブログです。

Pythonでクラス内の複数メソッドに一括でデコレータを当てる。

調べて試して、盛大に時間を使って詰まったので備忘録として残します。 本記事はPython3.6.4を想定しています。

背景

  • アプリケーションコード内には多くのクラスが定義されている。
  • それぞれのクラスには、またまた多くのメソッドが定義されている。
  • それら多くのメソッドの前後に、共通処理を挟みたくなった。
class A:
    def method1():
        ....

    def method2():
        ....

    ....(メソッドいっぱい)

class B:
    ....(メソッドいっぱい)

class C:
    ....(メソッドいっぱい)

....(クラスいっぱい)

やりたかったこと

  • クラスにデコレータを付けることで、一括で複数メソッドにデコレータを当てたい。
  • メソッドによってはデコレータを付けたくない。
  • 管理の点から、できるだけシンプルに書きたい。

完成メージ

@decorate_cls()
class A:
    def method1():
        ....

    def method2():
        ....

    ....

@decorate_cls()
class B:
    ....

@decorate_cls()
class C:
    ....

....

やりたくなかったこと

  • それぞれのメソッドにひとつずつデコレータを付ける。

結論

  • いろいろ試してみた結果、下記のような感じでデコレータを作成できた。
  • 登場人物は3つ
    • メソッドに当てる引数付きデコレータ: decorate_fn
    • クラスに当てる引数付きデコレータ: decorate_cls
    • デコレータを当てられるクラス: Test

サンプルコード

from functools import wraps
import inspect

# メソッドに当てる引数付きデコレータ
def decorate_fn(name=''):
    def wrapper(fn):
        @wraps(fn)
        def decorate(*args, **kwargs):
            print(f'----{name} start----')
            result = fn(*args, **kwargs)
            print(f'----{name} end----')
            return result
        return decorate
    return wrapper


# クラスに当てる引数付きデコレータ
def decorate_cls(exclude=[]):
    def decorate(Cls):
        for name, fn in inspect.getmembers(Cls):
            if name.startswith('__'):
                continue
            if callable(getattr(Cls, name)) and not name in exclude:
                setattr(Cls, name, decorate_fn(name)(fn))
        return Cls

    return decorate


# デコレータを当てられるクラス
@decorate_cls(exclude=['not_decorated'])
class Test():
    @classmethod
    def decorated_classmethod(cls):
        print('This is classmethod.')

    def decorated_method(self):
        print('This is decorated method.')

    def not_decorated(self):
        print('This is not decorated.')

if __name__ == '__main__':
    Test.decorated_classmethod()
    test = Test()
    test.decorated_method()
    test.not_decorated()

実行結果

----decorated_classmethod start----
This is classmethod.
----decorated_classmethod end----
----decorated_method start----
This is decorated method.
----decorated_method end----
This is not decorated.

解説

メソッドに当てる引数付きデコレータ

まずは def decorate_fn() から。

これは、メソッドに当てる引数付きデコレータ。 今回は引数としてメソッド名を受け、メソッドの実行前後に下記のようなprint文を挟んだ。

----<method name> start----
----<method name> end----

Python3における引数付きデコレータの書き方については、すでに良い記事があるのでそちらを参考にしてほしい。 qiita.com

自分はちゃんとfunctools.wrapsを使っています。

クラスに当てる引数付きデコレータ

次に、def decorate_cls()

とりあえず再掲。

# クラスに当てる引数付きデコレータ
def decorate_cls(exclude=[]):
    def decorate(Cls):
        for name, fn in inspect.getmembers(Cls):
            if name.startswith('__'):
                continue
            if callable(getattr(Cls, name)) and not name in exclude:
                setattr(Cls, name, decorate_fn(name)(fn))
        return Cls

    return decorate

中身の関数は、クラスオブジェクトを受けて全メソッドを取得し、特殊メソッドでないメソッドに対して先に定義したdecorate_fn()を当てて、それをもとのクラスにsetattr()しています。

分かりにくいところとしてはdecorate_fn(name)(fn)でしょうか。 decorate_fn()は引数付きデコレータですので、まずdecorate_fn(name)までの部分でデコレータに引数だけを与えています。 そして返ってくるデコレータに対して(fn)でメソッドを渡しています。

また、このデコレータも引数付きデコレータにしたかったので、引数を渡せるように関数の入れ子にしています。

デコレータを当てられるクラス

最後に class Test() いくつか適当にメソッドを定義しているクラスです。

最後に

晴れてメソッド一括デコレータを作れました。 あとはひたすらクラスにこのデコレータを当てるだけ。