Author: Manuel Jacob <m...@manueljacob.de>
Branch: py3.6
Changeset: r94258:236031b989f8
Date: 2018-04-06 14:27 +0200
http://bitbucket.org/pypy/pypy/changeset/236031b989f8/

Log:    Implement the logic that calls the async generator finalizer hook.

diff --git a/pypy/interpreter/generator.py b/pypy/interpreter/generator.py
--- a/pypy/interpreter/generator.py
+++ b/pypy/interpreter/generator.py
@@ -3,7 +3,7 @@
 from pypy.interpreter.pyopcode import LoopBlock, SApplicationException, Yield
 from pypy.interpreter.pycode import CO_YIELD_INSIDE_TRY
 from pypy.interpreter.astcompiler import consts
-from rpython.rlib import jit, rgc
+from rpython.rlib import jit, rgc, rweakref
 from rpython.rlib.objectmodel import specialize
 from rpython.rlib.rarithmetic import r_uint
 
@@ -588,20 +588,52 @@
     KIND = "async generator"
     KIND_U = u"async_generator"
 
+    def __init__(self, frame, name=None, qualname=None):
+        self.hooks_inited = False
+        GeneratorOrCoroutine.__init__(self, frame, name, qualname)
+
+    def init_hooks(self):
+        if self.hooks_inited:
+            return
+        self.hooks_inited = True
+
+        self.w_finalizer = self.space.appexec([], '''():
+            import sys
+            hooks = sys.get_asyncgen_hooks()
+            return hooks.finalizer''')
+
+    def _finalize_(self):
+        if self.frame is not None and self.frame.lastblock is not None:
+            if self.w_finalizer is not self.space.w_None:
+                # XXX: this is a hack to resurrect the weakref that was cleared
+                # before running _finalize_()
+                if self.space.config.translation.rweakref:
+                    self.frame.f_generator_wref = rweakref.ref(self)
+                try:
+                    self.space.call_function(self.w_finalizer, self)
+                except OperationError as e:
+                    e.write_unraisable(self.space, "async generator finalizer")
+                return
+        GeneratorOrCoroutine._finalize_(self)
+
     def descr__aiter__(self):
         """Return an asynchronous iterator."""
         return self
 
     def descr__anext__(self):
+        self.init_hooks()
         return AsyncGenASend(self, self.space.w_None)
 
     def descr_asend(self, w_arg):
+        self.init_hooks()
         return AsyncGenASend(self, w_arg)
 
     def descr_athrow(self, w_type, w_val=None, w_tb=None):
+        self.init_hooks()
         return AsyncGenAThrow(self, w_type, w_val, w_tb)
 
     def descr_aclose(self):
+        self.init_hooks()
         return AsyncGenAThrow(self, None, None, None)
 
 
diff --git a/pypy/interpreter/test/test_coroutine.py 
b/pypy/interpreter/test/test_coroutine.py
--- a/pypy/interpreter/test/test_coroutine.py
+++ b/pypy/interpreter/test/test_coroutine.py
@@ -537,6 +537,56 @@
         assert state == 2
     """
 
+    def test_async_aclose_in_finalize_hook_await_in_finally(self): """
+        import gc
+        import sys
+        import types
+
+        @types.coroutine
+        def coro():
+            yield 'coro'
+
+        state = 0
+        async def ag():
+            nonlocal state
+            try:
+                yield
+            finally:
+                state = 1
+                await coro()
+                state = 2
+
+        async def run():
+            a = ag()
+            async for i in a:
+                break
+            del a
+            gc.collect()
+            gc.collect()
+            gc.collect()
+        a = run()
+
+        a2 = None
+        assert sys.get_asyncgen_hooks() == (None, None)
+        def _finalize(g):
+            nonlocal a2
+            a2 = g.aclose()
+        sys.set_asyncgen_hooks(finalizer=_finalize)
+        assert state == 0
+        try:
+            a.send(None)
+        except StopIteration:
+            pass
+        assert a2.send(None) == 'coro'
+        assert state == 1
+        try:
+            a2.send(None)
+        except StopIteration:
+            pass
+        assert state == 2
+        sys.set_asyncgen_hooks(None, None)
+    """
+
     def test_async_anext_close(self): """
         async def ag():
             yield 42
_______________________________________________
pypy-commit mailing list
pypy-commit@python.org
https://mail.python.org/mailman/listinfo/pypy-commit

Reply via email to