Pythonのデコレータを作って理解する

@hogeの意味がわからない

Pythonに触れていると次のようなコードを見ることがあります

@hoge
def fuga():
  return "piyo"

この@hogeがなにをしているのかよくわからなかったので調べてみました。

デコレータ

@hogeはデコレータと呼ばれていて、さきほどのコードは次のコードの糖衣構文です(参考)。

def fuga():
  return "piyo"

fuga = hoge(fuga)

つまり、次の二つのコードは等価です(実際にはだいたい等価(参考))。

@hoge
def fuga():
  return "piyo"

def fuga():
  return "piyo"

fuga = hoge(fuga)

Pythonでは関数は第一級オブジェクトとして扱われるので、関数を引数にとったり、返り値として返すことができます。

つまり、ここではfugaという関数をつくり、そのfugaを関数を引数にとるhogeに渡して、hogeから返ってきた関数をfugaに割り当てているわけです。

すべてのデコレータはこのような糖衣構文であると理解して問題ないと思います。

基本のデコレータをつくる

上の説明をもとに基本的なデコレータをつくってみます。

def decorator(function): # 関数を引数にとる

    def retfunc(*args): # うけとった関数を使いながらいろんな処理をする関数をつくる
        print("デコレートされたよ")
        return function(*args) # デコレートされる関数

    return retfunc # つくった関数を返す

@decorator
def fuga():
    return "piyo"

print( fuga() )

出力

デコレートされたよ
piyo

デコレータであるdecoratorは上で説明したように実際には関数です。

decoratorの内部で関数を作って、その関数で受け取った関数を使って何かして、作った関数を返します。

fugaに渡すべき引数は*argsで引き渡しています。

引数をとるデコレータ

次のような引数をとるデコレータもみかけるとおもいます。

@decorator("arg1")
def fuga():
  return "piyo"

どうやったら作れるでしょうか。

デコレータに渡す引数を複数にしてやれば作れるでしょうか?

def decorator(function, arg1): # デコレータに引数をとりたい

    def retfunc(*args):
        print("{}を受け取りました".format(arg1))
        print("デコレートされたよ")
        return function(*args)

    return retfunc

@decorator("arg1")
def fuga():
    return "piyo"

print( fuga() )

error

TypeError: decorator() takes exactly 2 arguments (1 given)

引数が足りないと怒られちゃいました。

デコレータが上で説明したような糖衣構文であることを考えると、このコードはインタプリタにはこんな感じに見えているはずです。

fuga = decorator("arg1")(fuga)

引数が足りないと怒られるのは当然ですね。

ということは、decorator("arg1")の部分でさきほどの基本のデコレータを作って返す形にすれば大丈夫そうです。

def decorator(arg1): # デコレータに引数をとりたい
    
    def real_decorator(function): # ここでベーシックなデコレータを作る
        
        def retfunc(*args):
            print("{}を受け取りました".format(arg1))
            print("デコレートされたよ")
            return function(*args)

        return retfunc
    
    return real_decorator

@decorator("arg1")
def fuga():
    return "piyo"

print( fuga() )

出力

arg1を受け取りました
デコレートされたよ
piyo

クラスを使ったデコレータ

引数をとるデコレータはfuga = decorator("arg1")(fuga)の形を作ってやればうまく動いてくれました。

ということは、decorator("arg1")で初期化したオブジェクトに、関数を渡してcallするというアプローチでも動いてくれるでしょうか?

つまり

obj = decorator("arg1")
fuga = obj(fuga)

というアプローチでも動くのでしょうか?

動きます。

class class_base_decorator:
    
    def __init__(self, arg1):
        self.arg1 = arg1
        print("{}で初期化".format(self.arg1))
        
    def __call__(self, function): # オブジェクトを関数のように呼ぶ
        
        def retfunc(*args):
            print("{}をつかってデコレートしています".format(self.arg1))
            print("デコレートされたよ")
            return function(*args)

        return retfunc
    
@class_base_decorator("arg1")
def fuga():
    return "piyo"

print( fuga() )

出力

arg1で初期化
arg1をつかってデコレートしています
デコレートされたよ
piyo

メモ化するデコレータを作ってみる

n番目のフィボナッチ数を求めるプログラムを、デコレータをつかってメモ化をしてみます。フィボナッチ数を求めるプログラムとそのメモ化についてはこちらなどを参照してください。

n番目のフィボナッチ数の定義

f(0) = 0,
f(1) = 1,
f(n) = f(n-1) + f(n-2) 　　　//  (n ≧ 2)

まずは素のフィボナッチ数を再帰で求める関数

def fib(n):
    if n == 0 :
        return 0
    if n == 1:
        return 1
    return fib(n-1) + fib(n-2)

これをメモ化していきます

class class_base_memoize:
    
    def __init__(self,function):
        self.memo = {} # メモ用の辞書を用意する
        self.function = function
        
    def __call__(self):
        
        def memoized(n):
            if n not in self.memo: # 辞書にメモされているか
                self.memo[n] = self.function(n) # されてなければメモする
            return self.memo[n] # メモの内容を返す

        return memoized

@class_base_memoize
def fib(n):
    if n == 0 :
        return 0
    if n == 1:
        return 1
    return fib(n-1) + fib(n-2)

メモ化できました。

上の例ではクラスをつかってメモ用の辞書を用意しましたが、Pythonにはクロージャがあって定義時の変数を覚えておいてくれるので(参照)、もっとシンプルに次のように書けます。

def memoize(function):

    memo = {} # memoizedは何度呼び出されても、関数定義時にあったmemoを見に行く

    def memoized(n):
        if n not in memo:
            memo[n] = function(n)
        return memo[n]

    return memoized

@memoize
def fib(n):
    if n == 0 :
        return 0
    if n == 1:
        return 1
    return fib(n-1) + fib(n-2)

参考

Python Tips：Pythonでクロージャを使いたい - Life with Python
用語集 — Python 3.6.3 ドキュメント
Pythonのデコレータを理解するための12Step - Qiita
Decorators With Arguments in Python – Scott Lobdell
Katie Silverio Decorators, unwrapped How do they work PyCon 2017