Python でメソッドとクラスメソッドの両方に対応するデコレータ

通常のデコレータはインスタンスメソッドを想定していますが、クラスメソッドやスタティックメソッド、さらには関数としての呼び出しにも対応させたい場合があります。

問題の背景

単純なデコレータはメソッドに適用すると self を正しく処理しますが、クラスメソッドやスタティックメソッドとの組み合わせで問題が起きることがあります。

from functools import wraps

def logger(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        print(f"Calling {func.__name__}")
        return func(*args, **kwargs)
    return wrapper

class MyClass:
    @logger
    @classmethod
    def class_method(cls):
        pass  # 順序によってはエラー

ディスクリプタを使った解決策

デコレータ自体をディスクリプタにすることで、どのような呼び出し方にも対応できます。

from functools import wraps

class UniversalDecorator:
    def __init__(self, func):
        self.func = func
        wraps(func)(self)
    
    def __call__(self, *args, **kwargs):
        print(f"Calling {self.func.__name__}")
        return self.func(*args, **kwargs)
    
    def __get__(self, obj, objtype=None):
        if obj is None:
            return self
        return BoundMethod(self, obj)

class BoundMethod:
    def __init__(self, decorator, obj):
        self.decorator = decorator
        self.obj = obj
    
    def __call__(self, *args, **kwargs):
        print(f"Calling {self.decorator.func.__name__}")
        return self.decorator.func(self.obj, *args, **kwargs)

より簡潔な実装

types.MethodType を活用する方法もあります。

import types
from functools import wraps

class MethodDecorator:
    def __init__(self, func):
        self.func = func
        wraps(func)(self)
    
    def __call__(self, *args, **kwargs):
        print(f"Calling {self.func.__name__}")
        return self.func(*args, **kwargs)
    
    def __get__(self, obj, objtype=None):
        if obj is None:
            return self
        return types.MethodType(self, obj)

class Example:
    @MethodDecorator
    def method(self):
        return "method"

関数とメソッド両方に対応

@MethodDecorator
def standalone():
    return "standalone"

e = Example()
print(e.method())    # メソッドとして動作
print(standalone())  # 関数として動作

この手法は、ロギング、認証、キャッシュなど汎用的なデコレータを作る際に有用です。