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() を定義したい。

1つの方法は再帰すること。実装は意外と素朴。呼び出し関係の維持はそのまま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

もう1つは 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の話

1つ目に関してわりとひどくって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

pytestのfixtureとcontextlib.contextmanagerでの例外の取り扱い方の違い

前回の記事でpytestのfixtureでもteardownが実行されることを確実にするにはtry-finallyで囲む必要があるという風に書いてしまっていた。

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

これが嘘という話

pytestでは正しくteardownが実行できている

例えば以下のようなコードをpytestで実行した時に、after部分のprintがちゃんと呼ばれる。

import pytest


@pytest.fixture
def ob():
    ob = object()
    print("")
    print("<<< before", id(ob))
    yield ob
    print(">>> after", id(ob))


def test_it(ob):
    print("**test", id(ob), "**")
    1 / 0

>>> after 4375605680 という出力が行われている。

pytest -s .
================================================= test session starts =================================================
platform darwin -- Python 3.5.2, pytest-3.0.2, py-1.4.31, pluggy-0.3.1
rootdir:$HOME/vboxshare/venvs/my3/sandbox/daily/20170722/example_pytest, inifile:
collected 1 items

02exception/test_it.py
<<< before 4375605680
**test 4375605680 **
F>>> after 4375605680


====================================================== FAILURES =======================================================
_______________________________________________________ test_it _______________________________________________________

ob = <object object at 0x104ce71b0>

    def test_it(ob):
        print("**test", id(ob), "**")
>       1 / 0
E       ZeroDivisionError: division by zero

02exception/test_it.py:15: ZeroDivisionError
============================================== 1 failed in 0.05 seconds ===============================================

contextlibの実装の詳細

contextlib.contextmanagerは内部で_GeneratorContextManagerというクラスが使われているのだけれど。このクラスの __exit__() の定義が期待していたものとは違っていた。以下のような定義(長いし難しいので全部読む必要は無い)。

class _GeneratorContextManager(ContextDecorator):
    def __exit__(self, type, value, traceback):
        if type is None:
            try:
                next(self.gen)
            except StopIteration:
                return
            else:
                raise RuntimeError("generator didn't stop")
        else:
            if value is None:
                # Need to force instantiation so we can reliably
                # tell if we get the same exception back
                value = type()
            try:
                self.gen.throw(type, value, traceback)
                raise RuntimeError("generator didn't stop after throw()")
            except StopIteration as exc:
                # Suppress StopIteration *unless* it's the same exception that
                # was passed to throw().  This prevents a StopIteration
                # raised inside the "with" statement from being suppressed.
                return exc is not value
            except RuntimeError as exc:
                # Likewise, avoid suppressing if a StopIteration exception
                # was passed to throw() and later wrapped into a RuntimeError
                # (see PEP 479).
                if exc.__cause__ is value:
                    return False
                raise
            except:
                # only re-raise if it's *not* the exception that was
                # passed to throw(), because __exit__() must not raise
                # an exception unless __exit__() itself failed.  But throw()
                # has to raise the exception to signal propagation, so this
                # fixes the impedance mismatch between the throw() protocol
                # and the __exit__() protocol.
                #
                if sys.exc_info()[1] is not value:
                    raise

context managerについて

実際のコードの内容に触れる前におさらいが必要かも。

context managerの機能についておさらいしておくと(詳しくはここ)、

  • context managerとして機能するオブジェクトは __enter__()__exit__(exc_type, exc_value, traceback) のmethodを持ったもの
  • 内部で例外が発生しない場合には、__exit__()の引数は全部None
  • 内部で例外が発生した場合に、__exit__() の戻り値がtruthyな値の場合には発生した例外が無視される

実際に読むべき重要なところ

読むのは例外が発生した時なので、最初のifはelse部分、valueも通常は入っているはず。まぁそんなわけで真面目に読むべきはここの部分。

                self.gen.throw(type, value, traceback)
                raise RuntimeError("generator didn't stop after throw()")

ここでself.genは渡された関数の生成するgeneratorが束縛されている(generator関数は呼び出す際に戻り値としてgenerator objectを返す)。 そしてgenerator.throw()は、ほとんど __exit__()と同じ引数をを取り、そのgeneratorの現在の位置情報を利用して例外を送出する。

例えば雑な例をあげてみる。以下の様なコードは以下のような出力を返す。

def gen():
    print("hai")
    yield 1
    print("hoi")


it = gen()
print(next(it))
it.throw(Exception, Exception("hmm"))

nextでyield 1までは実行された後に、throw()を呼ぶ(tracebackは渡していないけれど。たぶん良い感じのネストしたerror reportに使われるかexceptionの発生位置として使われる)。

hai
1
Traceback (most recent call last):
  File "qr_8838MIf.py", line 9, in <module>
    it.throw(Exception, Exception("hmm"))
  File "qr_8838MIf.py", line 3, in gen
    yield 1
Exception: hmm

pytestの方の実装

pytestの方はというとfixtures.pyとrunner.pyのあたりをみれば良いのだけれど。雑にnext()を呼んでいる。この辺(紹介は雑)

fixtures.py

def call_fixture_func(fixturefunc, request, kwargs):
    yieldctx = is_generator(fixturefunc)
    if yieldctx:
        it = fixturefunc(**kwargs)
        res = next(it)

        def teardown():
            try:
                next(it)
            except StopIteration:
                pass
            else:
                fail_fixturefunc(fixturefunc,
                    "yield_fixture function has more than one 'yield'")

        request.addfinalizer(teardown)
    else:
        res = fixturefunc(**kwargs)
    return res

runner.py

    def _callfinalizers(self, colitem):
        finalizers = self._finalizers.pop(colitem, None)
        exc = None
        while finalizers:
            fin = finalizers.pop()
            try:
                fin()
            except Exception:
                # XXX Only first exception will be seen by user,
                #     ideally all should be reported.
                if exc is None:
                    exc = sys.exc_info()
        if exc:
            py.builtin._reraise(*exc)

簡略化した話

詳細に立ち寄りすぎたので簡略化すると例えば以下の様なコードがあったとする(再掲)。

def ob():
    ob = object()
    print("")
    print("<<< before", id(ob))

    yield ob  # ここで例外が発生した場合の話

    print(">>> after", id(ob))

ここで generator.throw() が呼ばれれば当然yieldの位置までしか実行はされずに終わる。一方でnext()が呼ばれれば通常通りiteratorを進めるだけなので以降の処理が実行される(最終的にStopIterationが送出される)。

雑な対応方法

上のことから考えて雑な対応方法を考えると以下の様な感じになる。

  • なるべく元のcontextlibの実装を利用する
  • 最後まで動かし切るために next() は呼ぶ
class SafeContextManager(_GeneratorContextManager):
    def __exit__(self, type, value, traceback):
        try:
            next(self.gen)
        except StopIteration:
            if type is None:
                return
        else:
            if type is None:
                raise RuntimeError("generator didn't stop")
        return super().__exit__(type, value, traceback)

上の対応方法に従った形で自分用のcontxtlib.contextmanagerを作ってあげれば良い(safe?という修飾は良くない気もするけれど)。

  • どんなときでもnext()は呼ぶ。
  • type is None の場合には、元の実装と同じ形で処理をする
  • 最も末尾で元の__exit__()を呼ぶ。

(元の__exit__()が呼ばれたとき、type is Noneになることはありえない=next()が2度呼ばれるということはありえない)

from functools import wraps

def safecontextmanager(func):
    @wraps(func)
    def helper(*args, **kwds):
        return SafeContextManager(func, args, kwds)

    return helper

おしまい。