Merged revisions 72912,72920,72940 via svnmerge from

svn+ssh://pythondev@svn.python.org/python/trunk

........
  r72912 | benjamin.peterson | 2009-05-25 08:13:44 -0500 (Mon, 25 May 2009) | 5 lines

  add a SETUP_WITH opcode

  It speeds up the with statement and correctly looks up the special
  methods involved.
........
  r72920 | benjamin.peterson | 2009-05-25 15:12:57 -0500 (Mon, 25 May 2009) | 1 line

  take into account the fact that SETUP_WITH pushes a finally block
........
  r72940 | benjamin.peterson | 2009-05-26 07:49:59 -0500 (Tue, 26 May 2009) | 1 line

  teach the peepholer about SETUP_WITH
........

10 files changed, 105 insertions, 84 deletions

diff --git a/Doc/library/dis.rst b/Doc/library/dis.rst
index 8d4d956..341beb5 100644
--- a/Doc/library/dis.rst
+++ b/Doc/library/dis.rst
@@ -436,6 +436,18 @@ the stack so that it is available for further iterations of the loop.
    by ``CALL_FUNCTION`` to construct a class.
 
 
+.. opcode:: SETUP_WITH (delta)
+
+   This opcode performs several operations before a with block starts.  First,
+   it loads :meth:`~object.__exit__` from the context manager and pushes it onto
+   the stack for later use by :opcode:`WITH_CLEANUP`.  Then,
+   :meth:`~object.__enter__` is called, and a finally block pointing to *delta*
+   is pushed.  Finally, the result of calling the enter method is pushed onto
+   the stack.  The next opcode will either ignore it (:opcode:`POP_TOP`), or
+   store it in (a) variable(s) (:opcode:`STORE_FAST`, :opcode:`STORE_NAME`, or
+   :opcode:`UNPACK_SEQUENCE`).
+
+
 .. opcode:: WITH_CLEANUP ()
 
    Cleans up the stack when a :keyword:`with` statement block exits.  TOS is
diff --git a/Doc/reference/compound_stmts.rst b/Doc/reference/compound_stmts.rst
index 016ccbb..4444f18 100644
--- a/Doc/reference/compound_stmts.rst
+++ b/Doc/reference/compound_stmts.rst
@@ -354,6 +354,8 @@ The execution of the :keyword:`with` statement with one "item" proceeds as follo
 
 #. The context expression is evaluated to obtain a context manager.
 
+#. The context manager's :meth:`__exit__` is loaded for later use.
+
 #. The context manager's :meth:`__enter__` method is invoked.
 
 #. If a target was included in the :keyword:`with` statement, the return value
@@ -363,9 +365,9 @@ The execution of the :keyword:`with` statement with one "item" proceeds as follo
 
       The :keyword:`with` statement guarantees that if the :meth:`__enter__`
       method returns without an error, then :meth:`__exit__` will always be
-      called.  Thus, if an error occurs during the assignment to the target
-      list, it will be treated the same as an error occurring within the suite
-      would be.  See step 5 below.
+      called. Thus, if an error occurs during the assignment to the target list,
+      it will be treated the same as an error occurring within the suite would
+      be. See step 6 below.
 
 #. The suite is executed.
 
diff --git a/Include/opcode.h b/Include/opcode.h
index 2b1c59d..7ffa359 100644
--- a/Include/opcode.h
+++ b/Include/opcode.h
@@ -130,8 +130,10 @@ extern "C" {
 #define CALL_FUNCTION_KW           141	/* #args + (#kwargs<<8) */
 #define CALL_FUNCTION_VAR_KW       142	/* #args + (#kwargs<<8) */
 
+#define SETUP_WITH 143
+
 /* Support for opargs more than 16 bits long */
-#define EXTENDED_ARG  143
+#define EXTENDED_ARG  144
 
 #define LIST_APPEND     145
 #define SET_ADD         146
diff --git a/Lib/opcode.py b/Lib/opcode.py
index 86ebc5c5..e8cccc3 100644
--- a/Lib/opcode.py
+++ b/Lib/opcode.py
@@ -166,12 +166,14 @@ hasfree.append(137)
 def_op('CALL_FUNCTION_VAR', 140)     # #args + (#kwargs << 8)
 def_op('CALL_FUNCTION_KW', 141)      # #args + (#kwargs << 8)
 def_op('CALL_FUNCTION_VAR_KW', 142)  # #args + (#kwargs << 8)
-def_op('EXTENDED_ARG', 143)
-EXTENDED_ARG = 143
+
+jrel_op('SETUP_WITH', 143)
 
 def_op('LIST_APPEND', 145)
 def_op('SET_ADD', 146)
 def_op('MAP_ADD', 147)
 
+def_op('EXTENDED_ARG', 144)
+EXTENDED_ARG = 144
 
 del def_op, name_op, jrel_op, jabs_op
diff --git a/Lib/test/test_descr.py b/Lib/test/test_descr.py
index 5c50571..ac5ab68 100644
--- a/Lib/test/test_descr.py
+++ b/Lib/test/test_descr.py
@@ -1569,6 +1569,7 @@ order (MRO) for bases """
         def some_number(self_, key):
             self.assertEqual(key, "hi")
             return 4
+        def swallow(*args): pass
 
         # It would be nice to have every special method tested here, but I'm
         # only listing the ones I can remember outside of typeobject.c, since it
@@ -1584,11 +1585,8 @@ order (MRO) for bases """
              set(("__class__",)), {}),
             ("__subclasscheck__", do_issubclass, return_true,
              set(("__bases__",)), {}),
-            # These two fail because the compiler generates LOAD_ATTR to look
-            # them up.  We'd have to add a new opcode to fix this, and it's
-            # probably not worth it.
-            # ("__enter__", run_context, iden),
-            # ("__exit__", run_context, iden),
+            ("__enter__", run_context, iden, set(), {"__exit__" : swallow}),
+            ("__exit__", run_context, swallow, set(), {"__enter__" : iden}),
             ]
 
         class Checker(object):
diff --git a/Python/ceval.c b/Python/ceval.c
index b5b5c27..b689f3d 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -119,6 +119,7 @@ static int import_all_from(PyObject *, PyObject *);
 static void format_exc_check_arg(PyObject *, const char *, PyObject *);
 static PyObject * unicode_concatenate(PyObject *, PyObject *,
                                       PyFrameObject *, unsigned char *);
+static PyObject * special_lookup(PyObject *, char *, PyObject **);
 
 #define NAME_ERROR_MSG \
 	"name '%.200s' is not defined"
@@ -2455,6 +2456,33 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
 					   STACK_LEVEL());
 			DISPATCH();
 
+		TARGET(SETUP_WITH)
+		{
+			static PyObject *exit, *enter;
+			w = TOP();
+			x = special_lookup(w, "__exit__", &exit);
+			if (!x)
+				break;
+			SET_TOP(x);
+		        u = special_lookup(w, "__enter__", &enter);
+			Py_DECREF(w);
+			if (!u) {
+				x = NULL;
+				break;
+			}
+			x = PyObject_CallFunctionObjArgs(u, NULL);
+			Py_DECREF(u);
+			if (!x)
+				break;
+			/* Setup the finally block before pushing the result
+			   of __enter__ on the stack. */
+			PyFrame_BlockSetup(f, SETUP_FINALLY, INSTR_OFFSET() + oparg,
+					   STACK_LEVEL());
+
+			PUSH(x);
+			DISPATCH();
+		}
+
 		TARGET(WITH_CLEANUP)
 		{
 			/* At the top of the stack are 1-3 values indicating
@@ -2479,17 +2507,36 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
 			   should still be resumed.)
 			*/
 
-			PyObject *exit_func = POP();
+			PyObject *exit_func;
 			u = TOP();
 			if (u == Py_None) {
+				POP();
+				exit_func = TOP();
+				SET_TOP(u);
 				v = w = Py_None;
 			}
 			else if (PyLong_Check(u)) {
+				POP();
+				switch(PyLong_AsLong(u)) {
+				case WHY_RETURN:
+				case WHY_CONTINUE:
+					/* Retval in TOP. */
+					exit_func = SECOND();
+					SET_SECOND(TOP());
+					SET_TOP(u);
+					break;
+				default:
+					exit_func = TOP();
+					SET_TOP(u);
+					break;
+				}
 				u = v = w = Py_None;
 			}
 			else {
-				v = SECOND();
+			        v = SECOND();
 				w = THIRD();
+				exit_func = stack_pointer[-7];
+				stack_pointer[-7] = NULL;
 			}
 			/* XXX Not the fastest way to call it... */
 			x = PyObject_CallFunctionObjArgs(exit_func, u, v, w,
@@ -2509,11 +2556,7 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
 			else if (err > 0) {
 				err = 0;
 				/* There was an exception and a True return */
-				STACKADJ(-2);
-				SET_TOP(PyLong_FromLong((long) WHY_SILENCED));
-				Py_DECREF(u);
-				Py_DECREF(v);
-				Py_DECREF(w);
+				PUSH(PyLong_FromLong((long) WHY_SILENCED));
 			}
 			PREDICT(END_FINALLY);
 			break;
@@ -3194,6 +3237,19 @@ fail: /* Jump here from prelude on failure */
 }
 
 
+static PyObject *
+special_lookup(PyObject *o, char *meth, PyObject **cache)
+{
+	PyObject *res;
+	res = _PyObject_LookupSpecial(o, meth, cache);
+	if (res == NULL && !PyErr_Occurred()) {
+		PyErr_SetObject(PyExc_AttributeError, *cache);
+		return NULL;
+	}
+	return res;
+}
+
+
 /* Logic for the raise statement (too complicated for inlining).
    This *consumes* a reference count to each of its arguments. */
 static enum why_code
diff --git a/Python/compile.c b/Python/compile.c
index c78949d..490137f 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -761,6 +761,8 @@ opcode_stack_effect(int opcode, int oparg)
 			return -1;
 		case BREAK_LOOP:
 			return 0;
+	        case SETUP_WITH:
+			return 7;
 		case WITH_CLEANUP:
 			return -1; /* XXX Sometimes more */
 		case STORE_LOCALS:
@@ -3085,85 +3087,31 @@ expr_constant(expr_ty e)
 static int
 compiler_with(struct compiler *c, stmt_ty s)
 {
-    static identifier enter_attr, exit_attr;
     basicblock *block, *finally;
-    identifier tmpvalue = NULL, tmpexit = NULL;
 
     assert(s->kind == With_kind);
 
-    if (!enter_attr) {
-	enter_attr = PyUnicode_InternFromString("__enter__");
-	if (!enter_attr)
-	    return 0;
-    }
-    if (!exit_attr) {
-	exit_attr = PyUnicode_InternFromString("__exit__");
-	if (!exit_attr)
-	    return 0;
-    }
-
     block = compiler_new_block(c);
     finally = compiler_new_block(c);
     if (!block || !finally)
 	return 0;
 
-    if (s->v.With.optional_vars) {
-	/* Create a temporary variable to hold context.__enter__().
-	   We need to do this rather than preserving it on the stack
-	   because SETUP_FINALLY remembers the stack level.
-	   We need to do the assignment *inside* the try/finally
-	   so that context.__exit__() is called when the assignment
-	   fails.  But we need to call context.__enter__() *before*
-	   the try/finally so that if it fails we won't call
-	   context.__exit__().
-	*/
-	tmpvalue = compiler_new_tmpname(c);
-	if (tmpvalue == NULL)
-	    return 0;
-	PyArena_AddPyObject(c->c_arena, tmpvalue);
-    }
-	tmpexit = compiler_new_tmpname(c);
-	if (tmpexit == NULL)
-	    return 0;
-	PyArena_AddPyObject(c->c_arena, tmpexit);
-
     /* Evaluate EXPR */
     VISIT(c, expr, s->v.With.context_expr);
+    ADDOP_JREL(c, SETUP_WITH, finally);
 
-    /* Squirrel away context.__exit__ by stuffing it under context */
-    ADDOP(c, DUP_TOP);
-    ADDOP_O(c, LOAD_ATTR, exit_attr, names);
-	if (!compiler_nameop(c, tmpexit, Store))
-	    return 0;
-
-    /* Call context.__enter__() */
-    ADDOP_O(c, LOAD_ATTR, enter_attr, names);
-    ADDOP_I(c, CALL_FUNCTION, 0);
-
-    if (s->v.With.optional_vars) {
-	/* Store it in tmpvalue */
-	if (!compiler_nameop(c, tmpvalue, Store))
-	    return 0;
-    }
-    else {
-	/* Discard result from context.__enter__() */
-	ADDOP(c, POP_TOP);
-    }
-
-    /* Start the try block */
-    ADDOP_JREL(c, SETUP_FINALLY, finally);
-
+    /* SETUP_WITH pushes a finally block. */
     compiler_use_next_block(c, block);
     if (!compiler_push_fblock(c, FINALLY_TRY, block)) {
 	return 0;
     }
 
     if (s->v.With.optional_vars) {
-	/* Bind saved result of context.__enter__() to VAR */
-	if (!compiler_nameop(c, tmpvalue, Load) ||
-	    !compiler_nameop(c, tmpvalue, Del))
-	  return 0;
-	VISIT(c, expr, s->v.With.optional_vars);
+        VISIT(c, expr, s->v.With.optional_vars);
+    }
+    else {
+        /* Discard result from context.__enter__() */
+        ADDOP(c, POP_TOP);
     }
 
     /* BLOCK code */
@@ -3181,9 +3129,6 @@ compiler_with(struct compiler *c, stmt_ty s)
     /* Finally block starts; context.__exit__ is on the stack under
        the exception or return information. Just issue our magic
        opcode. */
-	if (!compiler_nameop(c, tmpexit, Load) ||
-		!compiler_nameop(c, tmpexit, Del))
-		return 0;
     ADDOP(c, WITH_CLEANUP);
 
     /* Finally block ends. */
diff --git a/Python/import.c b/Python/import.c
index bccb971..23dd7b4 100644
--- a/Python/import.c
+++ b/Python/import.c
@@ -89,9 +89,10 @@ typedef unsigned short mode_t;
 			   change LIST_APPEND and SET_ADD, add MAP_ADD)
        Python 3.1a0: 3150 (optimize conditional branches:
 			   introduce POP_JUMP_IF_FALSE and POP_JUMP_IF_TRUE)
+       Python 3.2a0: 3160 (add SETUP_WITH)
 */
-#define MAGIC (3150 | ((long)'\r'<<16) | ((long)'\n'<<24))
 
+#define MAGIC (3160 | ((long)'\r'<<16) | ((long)'\n'<<24))
 /* Magic word as global; note that _PyImport_Init() can change the
    value of this global to accommodate for alterations of how the
    compiler works which are enabled by command line switches. */
diff --git a/Python/opcode_targets.h b/Python/opcode_targets.h
index 043e42a..deaf0a3 100644
--- a/Python/opcode_targets.h
+++ b/Python/opcode_targets.h
@@ -142,8 +142,8 @@ static void *opcode_targets[256] = {
 	&&TARGET_CALL_FUNCTION_VAR,
 	&&TARGET_CALL_FUNCTION_KW,
 	&&TARGET_CALL_FUNCTION_VAR_KW,
+	&&TARGET_SETUP_WITH,
 	&&TARGET_EXTENDED_ARG,
-	&&_unknown_opcode,
 	&&TARGET_LIST_APPEND,
 	&&TARGET_SET_ADD,
 	&&TARGET_MAP_ADD,
diff --git a/Python/peephole.c b/Python/peephole.c
index de1b2ac..23735b0 100644
--- a/Python/peephole.c
+++ b/Python/peephole.c
@@ -251,6 +251,7 @@ markblocks(unsigned char *code, Py_ssize_t len)
 			case SETUP_LOOP:
 			case SETUP_EXCEPT:
 			case SETUP_FINALLY:
+			case SETUP_WITH:
 				j = GETJUMPTGT(code, i);
 				blocks[j] = 1;
 				break;
@@ -566,6 +567,7 @@ PyCode_Optimize(PyObject *code, PyObject* consts, PyObject *names,
 			case SETUP_LOOP:
 			case SETUP_EXCEPT:
 			case SETUP_FINALLY:
+			case SETUP_WITH:
 				tgt = GETJUMPTGT(codestr, i);
 				/* Replace JUMP_* to a RETURN into just a RETURN */
 				if (UNCONDITIONAL_JUMP(opcode) &&
@@ -648,6 +650,7 @@ PyCode_Optimize(PyObject *code, PyObject* consts, PyObject *names,
 			case SETUP_LOOP:
 			case SETUP_EXCEPT:
 			case SETUP_FINALLY:
+			case SETUP_WITH:
 				j = addrmap[GETARG(codestr, i) + i + 3] - addrmap[i] - 3;
 				SETARG(codestr, i, j);
 				break;