pytestのyieldで指定するfixtureのやつを真似して実装してみた

github.com

pythonのyieldで指定するfixtureのやつを真似して実装してみた。意外と簡単にできるし。面白いし。便利な気がする。

yield fixture?

pytestのfixtureはsetup/teardownをgeneratorで定義できる。昔はyield_fixtureという名前だったけれど。現在はfixture自体がその機能を保持するようになった。

pytest supports execution of fixture specific finalization code when the fixture goes out of scope. By using a yield statement instead of return, all the code after the yield statement serves as the teardown code

pytest fixtures: explicit, modular, scalable — pytest documentation

fixture((テストの)実行に必要となる値)の後始末も含めてgeneratorの形式で記述できるので便利。

smtpのclientを終わったらcloseしている。

import smtplib
import pytest

@pytest.fixture(scope="module")
def smtp():
    smtp = smtplib.SMTP("smtp.gmail.com", 587, timeout=5)
    yield smtp  # provide the fixture value
    print("teardown smtp")
    smtp.close()

こういう風に書いてあげると以下の様な形で呼んでくれる。yield部分が間に挟まる処理でpytestの場合はtest部分。

setup    -- smtp
  test0
  test1
  test2
teardown -- smtp

また、yieldに渡した値はtest関数で使えるようになる。

def test(smtp):
    do_something()

この機能に似たような機能を作ってみようという話。

作ったものの使い方の話

作ったものは完全にpytestのそれと互換性のあるものではないし。pytestのコードを実際に確認して作ったわけではないので同じ実装になっているというわけではないけれど。

以下の様にして使う。

from yieldfixture import create
run, yield_fixture = create()


@yield_fixture
def f():
    print(">>> f")
    yield 1
    print(">>> f")


@yield_fixture
def g():
    print("  >>> g")
    yield 2
    print("  >>> g")


@run
def use_it(x, y):
    print("{} + {} = {}".format(x, y, x + y))

この実行結果は以下のようになる。

>>> f
  >>> g
1 + 2 = 3
  >>> g
>>> f

f,gとfixtureとして指定したのでその順に呼ばれている。runデコレータをつけた関数は実行されるので注意。 特に指定がない限りyield_fixtureデコレータをつけた関数はfixtureとして登録(指定)される。 yieldで渡された値が関数の引数として渡される。

(pytestのfixtureと違ってscopeオプションには対応していない。pytestの場合はscope=“module"などとするとそのmodule内で一度だけ実行みたいなことができる)

context

実行結果を見て分かる通り、最後の1 + 2 = 3という部分だけインデントされていない。この部分でもインデントをつけたい。そのためには何らかの方法で情報を受け渡せる必要がある。このような時にcontextを使う。

contextはargsとkwargsを持ったもの。positional argumentsとkeyword argumentsは適宜使い分けられる。

ctx = Context([0,1,2,3], {"foo": "bar"})

ctx[3] # => ctx.args[3]
ctx["foo"] # => ctx.kwargs["foo"]

これを利用するには with_context デコレータを一緒につける。残念ながらデコレータをつける順序を気にする必要があるので注意。

from yieldfixture import create, with_context
run, yield_fixture = create()


@yield_fixture
@with_context
def f(ctx):
    i = ctx["i"] = 0
    print("{}>>> f".format("  " * i))
    yield 1
    print("{}>>> f".format("  " * i))


@yield_fixture
@with_context
def g(ctx):
    i = ctx["i"] = ctx["i"] + 1
    print("{}>>> g".format("  " * i))
    yield 2
    print("{}>>> g".format("  " * i))


@run
def use_it(x, y, *, i=0):
    print("{}{} + {} = {}".format("  " * (i + 1), x, y, x + y))

contextに渡された値は関数の引数にキーワード引数として渡される。そんなわけで今度はindentを記録して使う事が出来るようになった。

>>> f
  >>> g
    1 + 2 = 3
  >>> g
>>> f

途中で発生する例外に対して安全なfixtureの書き方

特にtestでの利用を想定して作ったわけではないけれど。test時のsetup/teardownのことを考えると、途中のコードが失敗した時にも必ずteardown部分が実行されて欲しいということがあるかもしれない。この場合にはtry-finallyで括ってあげるのが無難(たぶんpytestの方もそうなはずだけれど確認していない(追記: 嘘でしたpytestの方try-finallyで囲まなくても大丈夫です))。

追記: 以下の部分のtry-finallyは不要です。0.2.0から不要になりました(詳細はコメント欄のリンク先の記事参照)。

from yieldfixture import create, with_context
run, yield_fixture = create()


@yield_fixture
@with_context
def f(ctx):
    i = ctx["i"] = 0
    print("{}>>> f".format(" " * i))
    try:
        yield 1
    finally:
        print("{}>>> f".format(" " * i))


@yield_fixture
@with_context
def g(ctx):
    i = ctx["i"] = ctx["i"] + 1
    print("{}>>> g".format(" " * i))
    try:
        yield 2
    finally:
        print("{}>>> g".format(" " * i))


@run
def use_it(x, y, *, i=0):
    print("{}{} + {} = {}".format(" " * (i + 1), x, y, x + y))
    1 / 0

途中でエラーになってもteardown部分(yield以降のコード)が実行される。

>>> f
 >>> g
  1 + 2 = 3
 >>> g
>>> f
Traceback (most recent call last):
  File "qr_88388ez.py", line 28, in <module>
    def use_it(x, y, *, i=0):
  File "$HOME/vboxshare/venvs/my3/sandbox/daily/20170722/example_yield_fixture/yieldfixture.py", line 98, in run_with
    return fn(*ctx.args, **ctx.kwargs)
  File "qr_88388ez.py", line 30, in use_it
    1 / 0
ZeroDivisionError: division by zero

利用するfixtureを直接指定する

利用するfixtureを直接指定することもできる。個人的にはこれができたのがかなり嬉しい感じだった。例えば、通常ならf,gと実行されるところをg,fと実行されるようにしたり、あるいはfだけを実行したりみたいなことが出来る。

from yieldfixture import create, with_context
run, yield_fixture = create()


@yield_fixture
@with_context
def f(ctx):
    i = ctx["i"] = ctx.get("i", -1) + 1
    print("{}>>> f".format("  " * i))
    try:
        yield 1
    finally:
        print("{}>>> f".format("  " * i))


@yield_fixture
@with_context
def g(ctx):
    i = ctx["i"] = ctx.get("i", -1) + 1
    print("{}>>> g".format("  " * i))
    try:
        yield 2
    finally:
        print("{}>>> g".format("  " * i))


@run
def use_it(x, y, *, i=0):
    print("{}{} + {} = {}".format("  " * (i + 1), x, y, x + y))


@run([g, f])
def use_it2(x, y, *, i=0):
    print("{}{} + {} = {}".format("  " * (i + 1), x, y, x + y))

2つ目のrunではfixtureの実行の順序が異なる。

>>> f
  >>> g
    1 + 2 = 3
  >>> g
>>> f
>>> g
  >>> f
    2 + 1 = 3
  >>> f
>>> g

実装の話

この機能自体の実装はそれほどたいへんではなかったりする。ちょっとだけ話したいので実装について話をすると、この機能は、任意長のcontext managerを順々にネストさせて実行する方法がわかれば実装できる。

任意長のcontext managerと言うのはこういう話。

例えば、f,g,hという順にfixtureが登録されたとして以下の様に実行されて欲しい。

setup -- f
  setup -- g
    setup -- h
      do something
    teardown -- h
  teardown -- g
teardown -- f

1つだけで良いのならcontextlib.contextmanager

ここで f にだけ考えると、contextlib.contextmanagerを使ってそのまま実装できる。

import contextlib

@contextlib.contextmanager
def f():
    # setup
    print("do something for setup")
    yield
    # teardown
    print("do something for teardown")

関数を自動的にwithで使えるcontext managerにしてくれるので便利。

with f():
    print("hello")


# do something for setup
# hello
# do something for teardown

context manager自体は、例えばrequestを投げて返ってきたresponseをcloseしたいみたいな時にも使われる。

@contextlib.contextmanager
def closing(response):
    try:
        yield response
    finally:
        response.close()

class Response:
    def __init__(self):
        print("open")

    def close(self):
        print("close")


with closing(Response()) as response:
    print(response)

# open
# <__main__.Response object at 0x105e54128>
# close

ちなみに同様の機能を持った関数は contextlib.closing() として存在していたりはする。

複数の任意個のcontext managerを扱うには?

複数の任意個のcontext managerを扱うにはどうするかというのが実装に関しては本題。

例えば、f,g,hというcontext managerがあったとしてこれを愚直に手で書くのならネストした記述も問題ない。

with f():
    with g():
        with h():
            print("do something")

これを例えば nested([f, g, h]) などという感じの呼び出しで実現する nested() を定義したい。

１つの方法は再帰すること。実装は意外と素朴。呼び出し関係の維持はそのままwith syntaxがやってくれる。

@contextlib.contextmanager
def nested(cms):
    if not cms:
        yield
    else:
        with cms[0]():
            with nested(cms[1:]):
                yield

ネストしたcontext managerの呼び出しの挙動が分かればまぁ直感的な気がする。

例えば以下の様なコードもしっかりネストされる(setupが全部呼び終わった後にteardownが逆順で呼ばれる)。

import contextlib


@contextlib.contextmanager
def f():
    # setup
    print("setup f")
    yield
    # teardown
    print("teardown f")


@contextlib.contextmanager
def g():
    # setup
    print("setup g")
    yield
    # teardown
    print("teardown g")


@contextlib.contextmanager
def h():
    # setup
    print("setup h")
    yield
    # teardown
    print("teardown h")


with nested([f, g, h]):
    print("** do something**")

# setup f
# setup g
# setup h
# ** do something**
# teardown h
# teardown g
# teardown f

もう１つは contextlib.ExitStack を使うこと。再帰など考えなくて良いので使う分にはこちらのほうが良いのかもしれない。 (ちなみに今回作ったyieldfixtureでは直接再帰を書いている)

@contextlib.contextmanager
def nested(cms):
    with contextlib.ExitStack() as s:
        for cm in cms:
            s.enter_context(cm())
        yield

これも同様に動く。

with nested([f, g, h]):
    print("** do something**")

# setup f
# setup g
# setup h
# ** do something**
# teardown h
# teardown g
# teardown f

補足的な話

これは全く脇道にそれた話題で補足的な話しになるけれど。contextlib.ExitStack は goのdefer的な感覚で使えるかもしれない。callbackを渡せるのでdeferと同様の感覚で実行するクロージャや関数を指定してあげれば良い。

import contextlib


class Response:
    def __init__(self, name):
        self.name = name
        print("open", self.name)

    def close(self):
        print("close", self.name)


with contextlib.ExitStack() as s:
    response = Response("1")
    s.callback(response.close)
    response = Response("2")
    s.callback(response.close)

    print("use", response.name, response)
    print("use", response.name, response)

closeが渡された順序と逆順に実行されていく(これもdeferとおんなじ感じ)。

open 1
open 2
use 2 <__main__.Response object at 0x10c7c4748>
use 2 <__main__.Response object at 0x10c7c4748>
close 2
close 1

ところで例外が起きた時にはどうなるかというと。

with contextlib.ExitStack() as s:
    response = Response("1")
    s.callback(response.close)
    response = Response("2")
    s.callback(response.close)

    print("use", response.name, response)
    1 / 0  # ZeroDivisionError !!!!!!11
    print("use", response.name, response)

大丈夫。

open 1
open 2
use 2 <__main__.Response object at 0x104d83710>
close 2
close 1
Traceback (most recent call last):
  File "qr_8838JFK.py", line 20, in <module>
    1 / 0
ZeroDivisionError: division by zero

これとは関係なく思ったこと

これとは関係なくこのパッケージを作っている途中で思ったことを幾つか。

• unit test を書くのが面倒だったので integration testしか書いていない
• readmeにjinja2のtemplate(kamidana)を使うのが便利だった

integration testの話

１つ目に関してわりとひどくってintegration testといっていてもやっていることは出力を確認しているだけ。

import os.path
import unittest
import textwrap
import contextlib

@contextlib.contextmanager
def _capture():
    from io import StringIO
    buf = StringIO()

    with contextlib.redirect_stdout(buf):
        yield buf

class Tests(unittest.TestCase)
    def test_it__with_context(self):
        with _capture() as buf:
            with open(os.path.join(os.path.dirname(__file__), "../../examples/01usecontext.py")
                      ) as rf:
                exec(rf.read())

        expected = textwrap.dedent(
            """
        >>> f
          >>> g
            1 + 2 = 3
          >>> g
        >>> f
        """
        )
        self.assertEqual(buf.getvalue().strip(), expected.strip())

しかも内部ではexamplesのファイルを実行(exec)しているだけ。ひどい。ただ contextlib.redirect_stdout で標準出力を奪えるので出力結果を確認するようなテストを書くには便利かもしれない。

(ちなみにopenを使っている辺りでふしぎな位置に改行があるのはyapfでフォーマットを整えているせい(既にソースコードのフォーマットは人間が逐一気にして管理するものじゃない気がするので黙って受け入れている))

readmeの話

readmeは利用例(example)だけで十分

readmeにモチベーションとか機能とか書かれていても使い方がわかんないとどうしようもないし。APIドキュメントに価値をあまり見出していない人間なので。コードの利用例(example)以外に興味がない。そしてexampleを列挙できない程度に高機能なライブラリは真面目にドキュメントを作らないとだめと言うような認識(そしてそのような高機能なライブラリはほとんどライフワークになる程度に維持管理や保守がつらいのであんまり作りたくない)。

そんなわけで、readmeでは、exampleを列挙して終わりにしようというのが最近の結論なのだけれど。exampleを書くのがわりとだるい。そしてmarkdownやReST形式のファイルの中に逐一コード例をコピペするのも馬鹿馬鹿しいのでjinja2のテンプレートを使って生成することにした。

example書くの意外とだるい

特定のファイル(examples/以下のファイル)を読み込んだ結果を埋め込むというようなことができれば良い感じ。真面目にドキュメントを作るならsphinxを使うだろうし、sphinxではliteralincludeみたいなディレクティブが存在するので便利に外部のファイルを埋め込めるのだけれど。さすがにreadmeなので頑張りたくないという気持ちがあった。

実際に使ったのはkamidanaというjinja2-cli(j2-cli)の亜種みたいなものなのだけれど。やりたいことはほとんどどのようなテンプレートエンジンでも出来ると思う。

こんな感じでmisc以下にreadme.j2というファイルを作ってreadmeを生成している。ReSTの形式でreadmeを作っているのでindent(こちらはjinja2に組み込み)のフィルターが便利だった。

yieldfixture
========================================

.. image:: https://travis-ci.org/podhmo/yieldfixture.svg?branch=master
    :target: https://travis-ci.org/podhmo/yieldfixture

how to use
----------------------------------------

.. code-block:: python

  {{"examples/00simple.py"|read_from_file|indent(2, False)}}

output

.. code-block::

  {{"python examples/00simple.py"|read_from_command|indent(2, False)}}

コード例とコード例の実行結果をreadmeに埋め込んでいる。readmeの生成自体は以下のような単純なMakefileでやっている。

readme:
  kamidana -a kamidana.additionals.reader misc/readme.j2 | sed 's@${HOME}@$$HOME@g' > README.rst