[Python-Dev] Rewrite @contextlib.contextmanager in C

[Chris Angelico]

On Tue, Aug 9, 2016 at 7:14 AM, Wolfgang Maier
<wolfgang.maier at biologie.uni-freiburg.de> wrote:
> Right, I think a fairer comparison would be to:
>
> class ctx2:
>     def __enter__(self):
>         self.it = iter(self)
>         return next(self.it)
>
>     def __exit__(self, *args):
>         try:
>             next(self.it)
>         except StopIteration:
>             pass
>
>     def __iter__(self):
>         yield
>
> With this change alone the slowdown diminishes to ~ 1.7x for me. The rest is
> probably the extra overhead for being able to pass exceptions raised inside
> the with block back into the generator and such.

I played around with a few other variants to see where the slowdown
is. They all work out pretty much the same as the above; my two
examples are both used the same way as contextlib.contextmanager is,
but are restrictive on what you can do.

import timeit
import contextlib
import functools

class ctx1:
    def __enter__(self):
        pass
    def __exit__(self, *args):
        pass

@contextlib.contextmanager
def ctx2():
    yield

class SimplerContextManager:
    """Like contextlib._GeneratorContextManager but way simpler.

    * Doesn't reinstantiate itself - just reinvokes the generator
    * Doesn't allow yielded objects (returns self)
    * Lacks a lot of error checking. USE ONLY AS DIRECTED.
    """
    def __init__(self, func):
        self.func = func
        functools.update_wrapper(self, func)
    def __call__(self, *a, **kw):
        self.gen = self.func(*a, **kw)
        return self
    def __enter__(self):
        next(self.gen)
        return self
    def __exit__(self, type, value, traceback):
        if type is None:
            try: next(self.gen)
            except StopIteration: return
            else: raise RuntimeError("generator didn't stop")
        try: self.gen.throw(type, value, traceback)
        except StopIteration: return True
        # Assume any instance of the same exception type is a proper reraise
        # This is way simpler than contextmanager normally does, and costs us
        # the ability to detect exception handlers that coincidentally raise
        # the same type of error (eg "except ZeroDivisionError: print(1/0)").
        except type: return False

# Check that it actually behaves correctly
@SimplerContextManager
def ctxdemo():
    print("Before yield")
    try:
        yield 123
    except ZeroDivisionError:
        print("Absorbing 1/0")
        return
    finally:
        print("Finalizing")
    print("After yield (no exception)")

with ctxdemo() as val:
    print("1/0 =", 1/0)
with ctxdemo() as val:
    print("1/1 =", 1/1)
#with ctxdemo() as val:
#    print("1/q =", 1/q)

@SimplerContextManager
def ctx3():
    yield

class TooSimpleContextManager:
    """Now this time you've gone too far."""
    def __init__(self, func):
        self.func = func
    def __call__(self):
        self.gen = self.func()
        return self
    def __enter__(self):
        next(self.gen)
    def __exit__(self, type, value, traceback):
        try: next(self.gen)
        except StopIteration: pass

@TooSimpleContextManager
def ctx4():
    yield

class ctx5:
    def __enter__(self):
        self.it = iter(self)
        return next(self.it)

    def __exit__(self, *args):
        try:
            next(self.it)
        except StopIteration:
            pass

    def __iter__(self):
        yield

t1 = timeit.timeit("with ctx1(): pass", setup="from __main__ import ctx1")
print("%.3f secs" % t1)

for i in range(2, 6):
    t2 = timeit.timeit("with ctx2(): pass", setup="from __main__
import ctx%d as ctx2"%i)
    print("%.3f secs" % t2)
    print("slowdown: -%.2fx" % (t2 / t1))


My numbers are:

0.320 secs
1.354 secs
slowdown: -4.23x
0.899 secs
slowdown: -2.81x
0.831 secs
slowdown: -2.60x
0.868 secs
slowdown: -2.71x

So compared to the tight custom-written context manager class, all the
"pass it a generator function" varieties look pretty much the same.
The existing contextmanager factory has several levels of indirection,
and that's probably where the rest of the performance difference comes
from, but there is some cost to the simplicity of the gen-func
approach.

My guess is that a C-implemented version could replace piles of
error-handling code with simple pointer comparisons (hence my
elimination of it), and may or may not be able to remove some of the
indirection. I'd say it'd land about in the same range as the other
examples here. Is that worth it?

ChrisA