diff --git a/src/aspectlib/py3support.py b/src/aspectlib/py3support.py index db9f907..3a0d09f 100644 --- a/src/aspectlib/py3support.py +++ b/src/aspectlib/py3support.py @@ -20,38 +20,32 @@ def decorate_advising_generator_py3(advising_function, cutpoint_function, bind): assert isgeneratorfunction(cutpoint_function) def advising_generator_wrapper_py3(*args, **kwargs): - if bind: - advisor = advising_function(cutpoint_function, *args, **kwargs) - else: - advisor = advising_function(*args, **kwargs) - if not isgenerator(advisor): - raise ExpectedGenerator("advising_function %s did not return a generator." % advising_function) - try: + gen = cutpoint_function(*args, **kwargs) + while True: + if bind: + advisor = advising_function(cutpoint_function, *args, **kwargs) + else: + advisor = advising_function(*args, **kwargs) + + if not isgenerator(advisor): + raise ExpectedGenerator("advising_function %s did not return a generator." % advising_function) + advice = next(advisor) - while True: - logdebug('Got advice %r from %s', advice, advising_function) - if advice is Proceed or advice is None or isinstance(advice, Proceed): - if isinstance(advice, Proceed): - args = advice.args - kwargs = advice.kwargs - gen = cutpoint_function(*args, **kwargs) - try: - result = yield from gen - except BaseException: - advice = advisor.throw(*sys.exc_info()) - else: - try: - advice = advisor.send(result) - except StopIteration: - return result - finally: - gen.close() - elif advice is Return: - return - elif isinstance(advice, Return): - raise StopIteration(advice.value) - else: - raise UnacceptableAdvice("Unknown advice %s" % advice) - finally: - advisor.close() + logdebug('Got advice %r from %s', advice, advising_function) + result = None + if advice is Proceed or advice is None or isinstance(advice, Proceed): + try: + result = next(gen) + advice = advisor.send(result) + except StopIteration as e: + return e.value + except BaseException: + advice = advisor.throw(*sys.exc_info()) + + if advice is Return: + yield result + elif isinstance(advice, Return): + yield advice.value + else: + raise UnacceptableAdvice("Unknown advice %s" % advice) return mimic(advising_generator_wrapper_py3, cutpoint_function)