Properly check for consistency with the third argument of

compile() when compiling an AST node.

Include/Python-ast.h

@@ -501,5 +501,5 @@ keyword_ty _Py_keyword(identifier arg, expr_ty value, PyArena *arena);
 alias_ty _Py_alias(identifier name, identifier asname, PyArena *arena);
 
 PyObject* PyAST_mod2obj(mod_ty t);
-mod_ty PyAST_obj2mod(PyObject* ast, PyArena* arena);
+mod_ty PyAST_obj2mod(PyObject* ast, PyArena* arena, int mode);
 int PyAST_Check(PyObject* obj);

Lib/test/test_compile.py

@@ -441,6 +441,20 @@ if 1:
             self.assert_(type(ast) == _ast.Module)
             co2 = compile(ast, '%s3' % fname, 'exec')
             self.assertEqual(co1, co2)
+            # the code object's filename comes from the second compilation step
+            self.assertEqual(co2.co_filename, '%s3' % fname)
+
+        # raise exception when node type doesn't match with compile mode
+        co1 = compile('print 1', '<string>', 'exec', _ast.PyCF_ONLY_AST)
+        self.assertRaises(TypeError, compile, co1, '<ast>', 'eval')
+
+        # raise exception when node type is no start node
+        self.assertRaises(TypeError, compile, _ast.If(), '<ast>', 'exec')
+
+        # raise exception when node has invalid children
+        ast = _ast.Module()
+        ast.body = [_ast.BoolOp()]
+        self.assertRaises(TypeError, compile, ast, '<ast>', 'exec')
 
 
 def test_main():

Parser/asdl_c.py

@@ -954,13 +954,20 @@ PyObject* PyAST_mod2obj(mod_ty t)
     return ast2obj_mod(t);
 }
 
-mod_ty PyAST_obj2mod(PyObject* ast, PyArena* arena)
+/* mode is 0 for "exec", 1 for "eval" and 2 for "single" input */
+mod_ty PyAST_obj2mod(PyObject* ast, PyArena* arena, int mode)
 {
     mod_ty res;
+    PyObject *req_type[] = {(PyObject*)Module_type, (PyObject*)Expression_type,
+                            (PyObject*)Interactive_type};
+    char *req_name[] = {"Module", "Expression", "Interactive"};
+    assert(0 <= mode && mode <= 2);
+
     init_types();
-    if (!PyObject_IsInstance(ast, (PyObject*)mod_type)) {
-        PyErr_SetString(PyExc_TypeError, "expected either Module, Interactive "
-                        "or Expression node");
+
+    if (!PyObject_IsInstance(ast, req_type[mode])) {
+        PyErr_Format(PyExc_TypeError, "expected %s node, got %.400s",
+                     req_name[mode], Py_TYPE(ast)->tp_name);
         return NULL;
     }
     if (obj2ast_mod(ast, &res, arena) != 0)
@@ -1016,7 +1023,7 @@ def main(srcfile):
                             )
         c.visit(mod)
         print >>f, "PyObject* PyAST_mod2obj(mod_ty t);"
-        print >>f, "mod_ty PyAST_obj2mod(PyObject* ast, PyArena* arena);"
+        print >>f, "mod_ty PyAST_obj2mod(PyObject* ast, PyArena* arena, int mode);"
         print >>f, "int PyAST_Check(PyObject* obj);"
         f.close()
 

Python/Python-ast.c

@@ -5944,13 +5944,20 @@ PyObject* PyAST_mod2obj(mod_ty t)
     return ast2obj_mod(t);
 }
 
-mod_ty PyAST_obj2mod(PyObject* ast, PyArena* arena)
+/* mode is 0 for "exec", 1 for "eval" and 2 for "single" input */
+mod_ty PyAST_obj2mod(PyObject* ast, PyArena* arena, int mode)
 {
     mod_ty res;
+    PyObject *req_type[] = {(PyObject*)Module_type, (PyObject*)Expression_type,
+                            (PyObject*)Interactive_type};
+    char *req_name[] = {"Module", "Expression", "Interactive"};
+    assert(0 <= mode && mode <= 2);
+
     init_types();
-    if (!PyObject_IsInstance(ast, (PyObject*)mod_type)) {
-        PyErr_SetString(PyExc_TypeError, "expected either Module, Interactive "
-                        "or Expression node");
+
+    if (!PyObject_IsInstance(ast, req_type[mode])) {
+        PyErr_Format(PyExc_TypeError, "expected %s node, got %.400s",
+                     req_name[mode], Py_TYPE(ast)->tp_name);
         return NULL;
     }
     if (obj2ast_mod(ast, &res, arena) != 0)

Python/bltinmodule.c

@@ -466,7 +466,7 @@ builtin_compile(PyObject *self, PyObject *args, PyObject *kwds)
 	char *str;
 	char *filename;
 	char *startstr;
-	int start;
+	int mode = -1;
 	int dont_inherit = 0;
 	int supplied_flags = 0;
 	PyCompilerFlags cf;
@@ -474,6 +474,7 @@ builtin_compile(PyObject *self, PyObject *args, PyObject *kwds)
 	Py_ssize_t length;
 	static char *kwlist[] = {"source", "filename", "mode", "flags",
 				 "dont_inherit", NULL};
+	int start[] = {Py_file_input, Py_eval_input, Py_single_input};
 
 	if (!PyArg_ParseTupleAndKeywords(args, kwds, "Oss|ii:compile",
 					 kwlist, &cmd, &filename, &startstr,
@@ -495,6 +496,18 @@ builtin_compile(PyObject *self, PyObject *args, PyObject *kwds)
 		PyEval_MergeCompilerFlags(&cf);
 	}
 
+	if (strcmp(startstr, "exec") == 0)
+		mode = 0;
+	else if (strcmp(startstr, "eval") == 0)
+		mode = 1;
+	else if (strcmp(startstr, "single") == 0)
+		mode = 2;
+	else {
+		PyErr_SetString(PyExc_ValueError,
+				"compile() arg 3 must be 'exec', 'eval' or 'single'");
+		return NULL;
+	}
+
 	if (PyAST_Check(cmd)) {
 		if (supplied_flags & PyCF_ONLY_AST) {
 			Py_INCREF(cmd);
@@ -505,7 +518,7 @@ builtin_compile(PyObject *self, PyObject *args, PyObject *kwds)
 			mod_ty mod;
 
 			arena = PyArena_New();
-			mod = PyAST_obj2mod(cmd, arena);
+			mod = PyAST_obj2mod(cmd, arena, mode);
 			if (mod == NULL) {
 				PyArena_Free(arena);
 				return NULL;
@@ -526,19 +539,6 @@ builtin_compile(PyObject *self, PyObject *args, PyObject *kwds)
 		cf.cf_flags |= PyCF_SOURCE_IS_UTF8;
 	}
 #endif
-	/* 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)
-		start = Py_eval_input;
-	else if (strcmp(startstr, "single") == 0)
-		start = Py_single_input;
-	else {
-		PyErr_SetString(PyExc_ValueError,
-				"compile() arg 3 must be 'exec'"
-				"or 'eval' or 'single'");
-		goto cleanup;
-	}
 
 	if (PyObject_AsReadBuffer(cmd, (const void **)&str, &length))
 		goto cleanup;
@@ -547,7 +547,7 @@ builtin_compile(PyObject *self, PyObject *args, PyObject *kwds)
 				"compile() expected string without null bytes");
 		goto cleanup;
 	}
-	result = Py_CompileStringFlags(str, filename, start, &cf);
+	result = Py_CompileStringFlags(str, filename, start[mode], &cf);
 cleanup:
 	Py_XDECREF(tmp);
 	return result;