[pypy-commit] pypy default: minimal support for greenlet.settrace()

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: 
Changeset: r80793:c7b77b0b686c
Date: 2015-11-20 15:08 +0100
http://bitbucket.org/pypy/pypy/changeset/c7b77b0b686c/

Log:	minimal support for greenlet.settrace()

diff --git a/lib_pypy/greenlet.py b/lib_pypy/greenlet.py
--- a/lib_pypy/greenlet.py
+++ b/lib_pypy/greenlet.py
@@ -88,9 +88,19 @@
         #
         try:
             unbound_method = getattr(_continulet, methodname)
+            _tls.leaving = current
             args, kwds = unbound_method(current, *baseargs, to=target)
-        finally:
             _tls.current = current
+        except:
+            _tls.current = current
+            if hasattr(_tls, 'trace'):
+                _run_trace_callback('throw')
+            _tls.leaving = None
+            raise
+        else:
+            if hasattr(_tls, 'trace'):
+                _run_trace_callback('switch')
+            _tls.leaving = None
         #
         if kwds:
             if args:
@@ -122,6 +132,34 @@
         return f.f_back.f_back.f_back   # go past start(), __switch(), switch()
 
 # ____________________________________________________________
+# Recent additions
+
+GREENLET_USE_GC = True
+GREENLET_USE_TRACING = True
+
+def gettrace():
+    return getattr(_tls, 'trace', None)
+
+def settrace(callback):
+    try:
+        prev = _tls.trace
+        del _tls.trace
+    except AttributeError:
+        prev = None
+    if callback is not None:
+        _tls.trace = callback
+    return prev
+
+def _run_trace_callback(event):
+    try:
+        _tls.trace(event, (_tls.leaving, _tls.current))
+    except:
+        # In case of exceptions trace function is removed
+        if hasattr(_tls, 'trace'):
+            del _tls.trace
+        raise
+
+# ____________________________________________________________
 # Internal stuff
 
 try:
@@ -143,22 +181,32 @@
     _tls.current = gmain
 
 def _greenlet_start(greenlet, args):
-    args, kwds = args
-    _tls.current = greenlet
     try:
-        res = greenlet.run(*args, **kwds)
-    except GreenletExit, e:
-        res = e
+        args, kwds = args
+        _tls.current = greenlet
+        try:
+            if hasattr(_tls, 'trace'):
+                _run_trace_callback('switch')
+            res = greenlet.run(*args, **kwds)
+        except GreenletExit, e:
+            res = e
+        finally:
+            _continuation.permute(greenlet, greenlet.parent)
+        return ((res,), None)
     finally:
-        _continuation.permute(greenlet, greenlet.parent)
-    return ((res,), None)
+        _tls.leaving = greenlet
 
 def _greenlet_throw(greenlet, exc, value, tb):
-    _tls.current = greenlet
     try:
-        raise exc, value, tb
-    except GreenletExit, e:
-        res = e
+        _tls.current = greenlet
+        try:
+            if hasattr(_tls, 'trace'):
+                _run_trace_callback('throw')
+            raise exc, value, tb
+        except GreenletExit, e:
+            res = e
+        finally:
+            _continuation.permute(greenlet, greenlet.parent)
+        return ((res,), None)
     finally:
-        _continuation.permute(greenlet, greenlet.parent)
-    return ((res,), None)
+        _tls.leaving = greenlet
diff --git a/pypy/module/test_lib_pypy/test_greenlet_tracing.py b/pypy/module/test_lib_pypy/test_greenlet_tracing.py
new file mode 100644
--- /dev/null
+++ b/pypy/module/test_lib_pypy/test_greenlet_tracing.py
@@ -0,0 +1,53 @@
+import py
+try:
+    from lib_pypy import greenlet
+except ImportError, e:
+    py.test.skip(e)
+
+class SomeError(Exception):
+    pass
+
+class TestTracing:
+    def test_greenlet_tracing(self):
+        main = greenlet.getcurrent()
+        actions = []
+        def trace(*args):
+            actions.append(args)
+        def dummy():
+            pass
+        def dummyexc():
+            raise SomeError()
+        oldtrace = greenlet.settrace(trace)
+        try:
+            g1 = greenlet.greenlet(dummy)
+            g1.switch()
+            g2 = greenlet.greenlet(dummyexc)
+            py.test.raises(SomeError, g2.switch)
+        finally:
+            greenlet.settrace(oldtrace)
+        assert actions == [
+            ('switch', (main, g1)),
+            ('switch', (g1, main)),
+            ('switch', (main, g2)),
+            ('throw', (g2, main)),
+        ]
+
+    def test_exception_disables_tracing(self):
+        main = greenlet.getcurrent()
+        actions = []
+        def trace(*args):
+            actions.append(args)
+            raise SomeError()
+        def dummy():
+            main.switch()
+        g = greenlet.greenlet(dummy)
+        g.switch()
+        oldtrace = greenlet.settrace(trace)
+        try:
+            py.test.raises(SomeError, g.switch)
+            assert greenlet.gettrace() is None
+        finally:
+            greenlet.settrace(oldtrace)
+        assert actions == [
+            ('switch', (main, g)),
+        ]