python のデコレータ関係が気になったので調べてみた

前回の記事で Flask アプリの unittest で redirect が想定通りできているかをテストする際は、
test_client の with ブロックの中で request.path を見ないとダメだよ
と、書いたのですが、redirect のテストのたびに毎回この with ブロックを書くのは面倒だなと思い
デコレーターでできないか試してみました。

意外と大変だったのでメモしておきます。

元コード

import unittest
from request import request

class TestFlaskBlog(unittest.TestCase):
    def setUp(self):
        """ テスト開始前の初期設定 """
        self.app = create_app(test_config={
            'TESTING':True,
        })
        self.client = self.app.test_client()

    def test_error_handle_404(self): 
        # 不明な url の場合はログインページにリダイレクトされる
        client = self.client
        with client: # ←これ毎回書くのダルくね？
             ret = client.get('/hoge', follow_redirects=True)
             assert request.path == '/login'

修正後のコード

import unittest
from request import request
from functools import wraps

class TestFlaskBlog(unittest.TestCase):
    def setUp(self):
        """ テスト開始前の初期設定 """
        self.app = create_app(test_config={
            'TESTING':True,
        })
        self.client = self.app.test_client()

    def with_client(self):
        def wrapper(test_func):
            @wraps(test_func)
            def inner(*args, **kwargs):
                with self.test_client():
                    return test_func(*args, **kwargs)
            return inner
        return wrapper

    @with_client
    def test_error_handle_404(self):
        # 不明な url の場合はログインページにリダイレクトされる
        ret = self.client.get('/hoge', follow_redirects=True)
        assert request.path == '/login'

デコレータとはなんぞや

正直今までデコレータの理解から逃げていたので今回はいい勉強になりました。
まず、デコレータの基本的な使い方。

def wrapper(func):
    def inner(*args, **kwargs):
         print("デコりました！")
         return func(*args, **kwargs)
    return inner

@wrapper
def func(a, b)
    print(a + b)

これで func(1, 2) を実行すると以下のような出力になります。

デコりました！
3

と、こんな感じで、複数の関数でお決まりの前処理のようなものがある場合に
デコレータが活躍します。

wraps ってなんぞや

これは前回紹介した本「ゼロからFlaskがよくわかる本」の中でしれっと使われていたのでが
何のためのものなのか調べてみました。

これをつけないと、デコレータをつけた関数名を参照しようとしたときに、
デコレータのついている関数ではなく、デコレータになっている関数の名前が表示されてしまうようです。

def wrapper(func):
    """ デコります """
    def inner(*args, **kwargs):
         print("デコりました！")
         return func(*args, **kwargs)
    return inner

@wrapper
def func(a, b)
    """ デコられます """
    print(a + b)

print(func.__name__)
print(func.__doc__)

<関数>.__name__ で関数名が取得できるはずですが、これを実行すると...

wrapper
デコります

となります。func 関数くんが自分を見失ってしまいました。
これを防ぐのが wraps になります。

def wrapper(func):
    """ デコります """
    @wraps(func)
    def inner(*args, **kwargs):
         print("デコりました！")
         return func(*args, **kwargs)
    return inner

@wrapper
def func(a, b)
    """ デコられます """
    print(a + b)

print(func.__name__)
print(func.__doc__)

実行結果

func
デコられます

ラッパーとなる関数で self を使いたい

デコレータの記法を使うには、ラッパーとなる関数には関数オブジェクトのみを渡すようにする必要があるようです。

そのため、ラッパーとなる関数のコアの部分は残しておいて、外側から追加の変数を渡す必要があります。

今回は、test_client がテスト開始時に作成され、それを self.client で保持しているため、 self をラッパー関数の引数に加える必要がありました。
そのため、最終的にラッパー関数は以下のようになりました。

(略)
    def with_client(self):
        def wrapper(test_func):
            @wraps(test_func)
            def inner(*args, **kwargs):
                with self.test_client():
                    return test_func(*args, **kwargs)
            return inner
        return wrapper
(略)

最後に

この記事書いてる途中で思ったけど、 redirect テスト用の共通関数作った方が早くね...？
まあ、勉強になったからヨシとします。