add a SETUP_WITH opcode

It speeds up the with statement and correctly looks up the special
methods involved.

3 files changed, 57 insertions, 57 deletions

diff --git a/Python/ceval.c b/Python/ceval.c
index 4f0877b9fa..92021e651a 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -128,6 +128,7 @@ static void format_exc_check_arg(PyObject *, char *, PyObject *);
 static PyObject * string_concatenate(PyObject *, PyObject *,
 				    PyFrameObject *, unsigned char *);
 static PyObject * kwd_as_string(PyObject *);
+static PyObject * special_lookup(PyObject *, char *, PyObject **);
 
 #define NAME_ERROR_MSG \
 	"name '%.200s' is not defined"
@@ -2467,6 +2468,33 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
 					   STACK_LEVEL());
 			continue;
 
+		case 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);
+			continue;
+		}
+
 		case WITH_CLEANUP:
 		{
 			/* At the top of the stack are 1-3 values indicating
@@ -3171,6 +3199,24 @@ fail: /* Jump here from prelude on failure */
 }
 
 
+static PyObject *
+special_lookup(PyObject *o, char *meth, PyObject **cache)
+{
+	PyObject *res;
+	if (PyInstance_Check(o)) {
+		if (!*cache)
+			return PyObject_GetAttrString(o, meth);
+		else
+			return PyObject_GetAttr(o, *cache);
+	}
+	res = _PyObject_LookupSpecial(o, meth, cache);
+	if (res == NULL && !PyErr_Occurred()) {
+		PyErr_SetObject(PyExc_AttributeError, *cache);
+		return NULL;
+	}
+	return res;
+}
+
 
 static PyObject *
 kwd_as_string(PyObject *kwd) {
diff --git a/Python/compile.c b/Python/compile.c
index 8c8530670b..e263d76342 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -778,6 +778,8 @@ opcode_stack_effect(int opcode, int oparg)
 			return -1;
 		case BREAK_LOOP:
 			return 0;
+	        case SETUP_WITH:
+			return 1;
 		case WITH_CLEANUP:
 			return -1; /* XXX Sometimes more */
 		case LOAD_LOCALS:
@@ -2821,80 +2823,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;
 
     assert(s->kind == With_kind);
 
-    if (!enter_attr) {
-	enter_attr = PyString_InternFromString("__enter__");
-	if (!enter_attr)
-	    return 0;
-    }
-    if (!exit_attr) {
-	exit_attr = PyString_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);
-    }
-
     /* 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);
-    ADDOP(c, ROT_TWO);
-
-    /* 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 */
diff --git a/Python/import.c b/Python/import.c
index 5e42e5c9f3..88aced06d1 100644
--- a/Python/import.c
+++ b/Python/import.c
@@ -74,9 +74,10 @@ typedef unsigned short mode_t;
        Python 2.7a0: 62171 (optimize list comprehensions/change LIST_APPEND)
        Python 2.7a0: 62181 (optimize conditional branches:
 			    introduce POP_JUMP_IF_FALSE and POP_JUMP_IF_TRUE)
+       Python 2.7a0  62191 (introduce SETUP_WITH)
 .
 */
-#define MAGIC (62181 | ((long)'\r'<<16) | ((long)'\n'<<24))
+#define MAGIC (62191 | ((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