1 files changed, 40 insertions, 13 deletions

diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
index b3d8b160e4..ccfce068b4 100644
--- a/Python/bltinmodule.c
+++ b/Python/bltinmodule.c
@@ -1,6 +1,7 @@
 /* Built-in functions */
 
 #include "Python.h"
+#include "Python-ast.h"
 
 #include "node.h"
 #include "code.h"
@@ -527,10 +528,43 @@ builtin_compile(PyObject *self, PyObject *args, PyObject *kwds)
 
 	cf.cf_flags = supplied_flags | PyCF_SOURCE_IS_UTF8;
 
-	str = source_as_string(cmd);
-	if (str == NULL)
+	if (supplied_flags &
+	    ~(PyCF_MASK | PyCF_MASK_OBSOLETE | PyCF_DONT_IMPLY_DEDENT | PyCF_ONLY_AST))
+	{
+		PyErr_SetString(PyExc_ValueError,
+				"compile(): unrecognised flags");
 		return NULL;
+	}
+	/* XXX Warn if (supplied_flags & PyCF_MASK_OBSOLETE) != 0? */
 
+	if (!dont_inherit) {
+		PyEval_MergeCompilerFlags(&cf);
+	}
+
+	if (PyAST_Check(cmd)) {
+		PyObject *result;
+		if (supplied_flags & PyCF_ONLY_AST) {
+			Py_INCREF(cmd);
+			result = cmd;
+		}
+		else {
+			PyArena *arena;
+			mod_ty mod;
+
+			arena = PyArena_New();
+			mod = PyAST_obj2mod(cmd, arena);
+			if (mod == NULL) {
+				PyArena_Free(arena);
+				return NULL;
+			}
+			result = (PyObject*)PyAST_Compile(mod, filename,
+							  &cf, arena);
+			PyArena_Free(arena);
+		}
+		return result;
+	}
+
+	/* XXX: is it possible to pass start to the PyAST_ branch? */
 	if (strcmp(startstr, "exec") == 0)
 		start = Py_file_input;
 	else if (strcmp(startstr, "eval") == 0)
@@ -539,22 +573,15 @@ builtin_compile(PyObject *self, PyObject *args, PyObject *kwds)
 		start = Py_single_input;
 	else {
 		PyErr_SetString(PyExc_ValueError,
-		   "compile() arg 3 must be 'exec' or 'eval' or 'single'");
+				"compile() arg 3 must be 'exec'"
+				"or 'eval' or 'single'");
 		return NULL;
 	}
 
-	if (supplied_flags &
-	    ~(PyCF_MASK | PyCF_MASK_OBSOLETE | PyCF_DONT_IMPLY_DEDENT | PyCF_ONLY_AST))
-	{
-		PyErr_SetString(PyExc_ValueError,
-				"compile(): unrecognised flags");
+	str = source_as_string(cmd);
+	if (str == NULL)
 		return NULL;
-	}
-	/* XXX Warn if (supplied_flags & PyCF_MASK_OBSOLETE) != 0? */
 
-	if (!dont_inherit) {
-		PyEval_MergeCompilerFlags(&cf);
-	}
 	return Py_CompileStringFlags(str, filename, start, &cf);
 }