[Jython-checkins] SF.net SVN: jython:[6538] trunk/jython/src/org/python/core/ContextGuard. java

[<th...@us...>]

Revision: 6538
          http://jython.svn.sourceforge.net/jython/?rev=6538&view=rev
Author:   thobes
Date:     2009-07-15 02:34:53 +0000 (Wed, 15 Jul 2009)

Log Message:
-----------
Added a tentative implementation of an optimized call path for using generators
as context managers (as contextlib.contextmanager does by wrapping generators).
This implementation unwraps the code object and wraps it in context manager
instead of wrapping the entire generator in a context manager.

Modified Paths:
--------------
    trunk/jython/src/org/python/core/ContextGuard.java

Modified: trunk/jython/src/org/python/core/ContextGuard.java
===================================================================
--- trunk/jython/src/org/python/core/ContextGuard.java	2009-07-15 02:31:42 UTC (rev 6537)
+++ trunk/jython/src/org/python/core/ContextGuard.java	2009-07-15 02:34:53 UTC (rev 6538)
@@ -1,10 +1,8 @@
 package org.python.core;
 
-/** Straightens the call path for some common cases
+/**
+ * Straightens the call path for some common cases
  */
-
-// XXX - add support for generators in conjunction w/ contextlib.contextmanager
-
 public class ContextGuard implements ContextManager {
 
     private final PyObject __enter__method;
@@ -19,7 +17,8 @@
         return __enter__method.__call__(ts);
     }
 
-    public boolean __exit__(ThreadState ts, PyObject type, PyObject value, PyObject traceback) {
+    public boolean __exit__(ThreadState ts, PyObject type, PyObject value,
+            PyObject traceback) {
         return __exit__method.__call__(ts, type, value, traceback).__nonzero__();
     }
 
@@ -30,4 +29,123 @@
             return new ContextGuard(manager);
         }
     }
+
+    // XXX - tentative support for generators in conjunction w/ contextlib.contextmanager
+    
+    /* Sample usage:
+     * 
+     * from org.python.core.ContextGuard import makeManager as contextmanager
+     * @contextmanager
+     * def my_manager():
+     *     print "setup"
+     *     try:
+     *         yield
+     *     finally:
+     *         print "done"
+     */
+    
+    public static PyObject makeManager(PyObject object) {
+        if (object instanceof PyFunction) {
+            PyFunction function = (PyFunction) object;
+            PyCode code = function.func_code;
+            if (code instanceof PyBaseCode) {
+                PyBaseCode pyCode = (PyBaseCode) code;
+                if (pyCode.co_flags.isFlagSet(CodeFlag.CO_GENERATOR)) {
+                    return new PyFunction(function.func_globals,
+                            function.func_defaults,
+                            new ContextCode(pyCode),
+                            function.__doc__,
+                            (function.func_closure == null) ? null : 
+                            ((PyTuple)function.func_closure).getArray());
+                }
+            }
+        }
+        throw Py.TypeError("Argument must be a generator function.");
+    }
+
+    private static class ContextCode extends PyBaseCode {
+        private final PyBaseCode code;
+        ContextCode(PyBaseCode code) {
+            this.co_name = code.co_name;
+            this.code = code;
+            this.co_argcount = code.co_argcount;
+            this.nargs = code.nargs;
+            this.co_firstlineno = code.co_firstlineno;
+            this.co_varnames = code.co_varnames;
+            this.co_cellvars = code.co_cellvars;
+            this.jy_npurecell = code.jy_npurecell;
+            this.co_freevars = code.co_freevars;
+            this.co_filename = code.co_filename;
+            this.co_nlocals = code.co_nlocals;
+            this.varargs = code.varargs;
+            this.varkwargs = code.varkwargs;
+            for (CodeFlag flag : CodeFlag.values()) {
+                if (code.co_flags.isFlagSet(flag) && flag != CodeFlag.CO_GENERATOR) {
+                    this.co_flags.setFlag(flag);
+                }
+            }
+        }
+        @SuppressWarnings("serial")
+        @Override
+        protected PyObject interpret(PyFrame frame, ThreadState ts) {
+            frame.f_back = null;
+            return new GeneratorContextManager(frame) {
+                @Override
+                PyObject body(ThreadState ts) {
+                    return code.interpret(frame, ts);
+                }
+            };
+        }
+        @SuppressWarnings("serial")
+        @Override
+        public PyObject call(ThreadState ts, PyFrame frame, final PyObject closure) {
+            frame.f_back = null;
+            return new GeneratorContextManager(frame) {
+                @Override
+                PyObject body(ThreadState ts) {
+                    return code.call(ts, frame, closure);
+                }
+            };
+        }
+    }
+    
+    @SuppressWarnings("serial")
+    private static abstract class GeneratorContextManager extends PyObject implements ContextManager {
+        final PyFrame frame;
+
+        public GeneratorContextManager(PyFrame frame) {
+            this.frame = frame;
+        }
+
+        public PyObject __enter__(ThreadState ts) {
+            PyObject res = body(ts);
+            if (frame.f_lasti == -1) {
+                throw Py.RuntimeError("generator didn't yield");
+            }
+            return res;
+        }
+        
+        public boolean __exit__(ThreadState ts, PyObject type, PyObject value,
+                PyObject traceback) {
+            PyException ex = null;
+            if (type != null && type != Py.None) {
+                ex = Py.makeException(type, value, traceback);
+                frame.setGeneratorInput(ex);
+            }
+            PyObject res = null;
+            try {
+                res = body(ts);
+            } catch(PyException e) {
+                if (e == ex) {
+                    return false;
+                }
+            }
+            if (frame.f_lasti != -1) {
+                throw Py.RuntimeError("generator didn't stop");
+            }
+            return res.__nonzero__();
+        }
+
+        abstract PyObject body(ThreadState ts);
+    }
 }


This was sent by the SourceForge.net collaborative development platform, the world's largest Open Source development site.