Patch #1810 by Thomas Lee, reviewed by myself:

allow compiling Python AST objects into code objects
in compile().

1 files changed, 44 insertions, 19 deletions

diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
index 1c36fe5..fff4fda 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"
@@ -481,6 +482,41 @@ builtin_compile(PyObject *self, PyObject *args, PyObject *kwds)
 
 	cf.cf_flags = supplied_flags;
 
+	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)) {
+		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;
+	}
+
 #ifdef Py_USING_UNICODE
 	if (PyUnicode_Check(cmd)) {
 		tmp = PyUnicode_AsUTF8String(cmd);
@@ -490,14 +526,7 @@ builtin_compile(PyObject *self, PyObject *args, PyObject *kwds)
 		cf.cf_flags |= PyCF_SOURCE_IS_UTF8;
 	}
 #endif
-	if (PyObject_AsReadBuffer(cmd, (const void **)&str, &length))
-		return NULL;
-	if ((size_t)length != strlen(str)) {
-		PyErr_SetString(PyExc_TypeError,
-				"compile() expected string without null bytes");
-		goto cleanup;
-	}
-
+	/* 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)
@@ -506,21 +535,17 @@ 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'");
 		goto cleanup;
 	}
 
-	if (supplied_flags &
-	    ~(PyCF_MASK | PyCF_MASK_OBSOLETE | PyCF_DONT_IMPLY_DEDENT | PyCF_ONLY_AST))
-	{
-		PyErr_SetString(PyExc_ValueError,
-				"compile(): unrecognised flags");
+	if (PyObject_AsReadBuffer(cmd, (const void **)&str, &length))
+		goto cleanup;
+	if ((size_t)length != strlen(str)) {
+		PyErr_SetString(PyExc_TypeError,
+				"compile() expected string without null bytes");
 		goto cleanup;
-	}
-	/* XXX Warn if (supplied_flags & PyCF_MASK_OBSOLETE) != 0? */
-
-	if (!dont_inherit) {
-		PyEval_MergeCompilerFlags(&cf);
 	}
 	result = Py_CompileStringFlags(str, filename, start, &cf);
 cleanup: