add a SETUP_WITH opcode

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

1 files changed, 46 insertions, 0 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) {