context_inside_generator_test.py

"""Testing generator that yields within a context and interrupt by error

Looks like the context is being closed by the python garbage collection if
error raised during the iteration.

https://stackoverflow.com/a/19825163/4145300

"""
import sys
from contextlib import contextmanager

PY3 = sys.version_info[0] == 3


class Context(object):

    def __init__(self):
        self.state = None

    def __enter__(self):
        print("->")
        self.state = "O"

    def __exit__(self, *args, **kwargs):
        print("<-")
        self.state = "X"


def generator(ctx):
    with ctx:
        for i in range(10):
            yield i


def run_and_err(arg):
    raise RuntimeError()


def case_a():
    print("\n*** Case A ***")
    context = Context()
    try:
        for i in generator(context):
            run_and_err(i)
    except Exception:
        pass

    print("State: %s" % context.state)


def case_b():
    print("\n*** Case B ***")
    context = Context()
    try:
        gen = generator(context)
        for i in gen:
            run_and_err(i)
    except Exception:
        pass

    print("State: %s" % context.state)


case_a()  # Context is closed by GC after for loop ends
case_b()  # Context is closed by GC after function ends


def fix_b():
    print("\n*** Fix B ***")
    context = Context()
    try:
        gen = generator(context)
        for i in gen:
            run_and_err(i)
    except RuntimeError:
        # Manually call `generator.close` so the context within
        # so the context will be closed together.
        gen.close()

    print("State: %s" % context.state)


fix_b()


if PY3:
    from contextlib import ExitStack

    class Context2(object):

        def __init__(self):
            self.state = None

        def __enter__(self):
            print("->")
            self.state = "O"
            with ExitStack() as cm:
                cm.callback(lambda x: x.close(), self)
            return self

        def __exit__(self, *args, **kwargs):
            print("<-")
            self.state = "X"

        close = __exit__

    def fix_b2():
        print("\n*** Fix B ***")
        context = Context2()
        try:
            gen = generator(context)
            for i in gen:
                run_and_err(i)
        except RuntimeError:
            pass

        print("State: %s" % context.state)

    fix_b2()


state = {"_": None}


@contextmanager
def ContextM():
    try:
        print("->")
        state["_"] = "O"
        yield
    finally:
        print("<-")
        state["_"] = "X"


def fix_b3():
    print("\n*** Fix B ***")
    try:
        context = ContextM()
        gen = generator(context)
        for i in gen:
            run_and_err(i)
    except RuntimeError:
        pass

    print("State: %s" % state["_"])


fix_b3()

Result:

*** Case A ***
->
<-
State: X

*** Case B ***
->
State: O
<-

*** Fix B ***
->
<-
State: X

*** Fix B ***
->
<-
State: X
<-

*** Fix B ***
->
State: O
<-