こんにちは、Pythonエンジニア見習いです。最近TemplateMethodパターンを使っているコードのリファクタリングをしていたところ、継承先にもデコレーターを適用したい場面に遭遇しました。その時、単純に継承元に適用するだけではうまくいかず、工夫が必要でした。今回はPythonでデコレーターを継承先にも適用する方法を皆様に紹介したいと思います。

なぜ単純に継承元に適応するだけではダメなのか

継承元の関数は継承先の関数はたとえ名前が同じであってもそれが指す関数オブジェクトは別物です。よって継承元でデコレーターを適用した場合、継承先ではそのデコレーターの適用された継承元の関数は継承先の同名の関数で上書きされます。最終的に継承先の関数ではデコレーターの適用されていない関数が残ります。

どうやったら継承先にデコレーターを適用できるのか

メンバ関数がデコレートされるタイミングはクラス(インスタンスではなく)が生成される時です。つまり、__new__を用いて、関数オブジェクト生成に介入すればうまくいくことがわかります。

例えば以下のようなコードの場合には

class A(metaclass=ABCMeta):

  @abc.abstractmethod
  def foo(self):
    return


class B(A):

  def foo(self):
    print("B")


class C(A):

  def foo(self):
    print("C")

クラスAを以下のようにすれば、その継承先のB、Cでもデコレーターdecoが適用されます。

class A(metaclass=ABCMeta):

  def __new__(cls, *args, **kwargs):
      cls.foo = deco(cls.foo)
      return super().__new__(cls)

  @abc.abstractmethod
  def foo(self):
    return

使用例

以下のコードは典型的なTemplateMethodパターンで組まれたコードに継承先にも適用されるデコレーターを使ったものです。また、このデコレーターは2つの関数を取り第一引数に通常実行される関数、第二引数に第一引数の関数の中で例外が発生した際に呼び出される関数を適用します。

import abc
import functools


class HogeException(Exception):
    pass


class FooException(HogeException):
    pass


class BarException(HogeException):
    pass

# fが通常実行される関数で、ffが例外発生時に呼ばれる関数
def deco(f, ff):
    @functools.wraps(f)
    def inner(*args, **kwargs):
        try:
            return f(*args, **kwargs)
        except HogeException:
            return ff(*args, **kwargs)
    return inner


class A(metaclass=abc.ABCMeta):

    @abc.abstractclassmethod
    def foo(self):
        pass

    # この関数のなかでデコレーターが適用されている
    def __new__(cls, *args, **kwargs):
        cls.foo = deco(cls.foo, cls.foo_sub)
        return super().__new__(cls)

    def foo_sub(self):
        print("foo_sub")


class B(A):

    def foo(self):
        raise FooException


class C(A):

    def foo(self):
        raise BarException


if __name__ == "__main__":
    # ちゃんと例外発生時に呼ばれる関数が呼ばれていることがわかる
    B().foo() # foo_sub
    C().foo() # foo_sub

まとめ

継承先にもデコレーターを適用したい場合は、__new__関数のなかでデコレーターを適用すればうまくいきます。