Skip to content

Commit f2bfd54

Browse files
birkenfeldbirkenfeld
committed
Properly check for consistency with the third argument of
compile() when compiling an AST node.
1 parent ea13dc6 commit f2bfd54

5 files changed

Lines changed: 54 additions & 26 deletions

File tree

‎Include/Python-ast.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -501,5 +501,5 @@ keyword_ty _Py_keyword(identifier arg, expr_ty value, PyArena *arena);
501501
alias_ty _Py_alias(identifier name, identifier asname, PyArena *arena);
502502

503503
PyObject* PyAST_mod2obj(mod_ty t);
504-
mod_ty PyAST_obj2mod(PyObject* ast, PyArena* arena);
504+
mod_ty PyAST_obj2mod(PyObject* ast, PyArena* arena, int mode);
505505
int PyAST_Check(PyObject* obj);

‎Lib/test/test_compile.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -441,6 +441,20 @@ def test_compile_ast(self):
441441
self.assert_(type(ast) == _ast.Module)
442442
co2 = compile(ast, '%s3' % fname, 'exec')
443443
self.assertEqual(co1, co2)
444+
# the code object's filename comes from the second compilation step
445+
self.assertEqual(co2.co_filename, '%s3' % fname)
446+
447+
# raise exception when node type doesn't match with compile mode
448+
co1 = compile('print 1', '<string>', 'exec', _ast.PyCF_ONLY_AST)
449+
self.assertRaises(TypeError, compile, co1, '<ast>', 'eval')
450+
451+
# raise exception when node type is no start node
452+
self.assertRaises(TypeError, compile, _ast.If(), '<ast>', 'exec')
453+
454+
# raise exception when node has invalid children
455+
ast = _ast.Module()
456+
ast.body = [_ast.BoolOp()]
457+
self.assertRaises(TypeError, compile, ast, '<ast>', 'exec')
444458

445459

446460
def test_main():

‎Parser/asdl_c.py‎

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -954,13 +954,20 @@ class PartingShots(StaticVisitor):
954954
return ast2obj_mod(t);
955955
}
956956
957-
mod_ty PyAST_obj2mod(PyObject* ast, PyArena* arena)
957+
/* mode is 0 for "exec", 1 for "eval" and 2 for "single" input */
958+
mod_ty PyAST_obj2mod(PyObject* ast, PyArena* arena, int mode)
958959
{
959960
mod_ty res;
961+
PyObject *req_type[] = {(PyObject*)Module_type, (PyObject*)Expression_type,
962+
(PyObject*)Interactive_type};
963+
char *req_name[] = {"Module", "Expression", "Interactive"};
964+
assert(0 <= mode && mode <= 2);
965+
960966
init_types();
961-
if (!PyObject_IsInstance(ast, (PyObject*)mod_type)) {
962-
PyErr_SetString(PyExc_TypeError, "expected either Module, Interactive "
963-
"or Expression node");
967+
968+
if (!PyObject_IsInstance(ast, req_type[mode])) {
969+
PyErr_Format(PyExc_TypeError, "expected %s node, got %.400s",
970+
req_name[mode], Py_TYPE(ast)->tp_name);
964971
return NULL;
965972
}
966973
if (obj2ast_mod(ast, &res, arena) != 0)
@@ -1016,7 +1023,7 @@ def main(srcfile):
10161023
)
10171024
c.visit(mod)
10181025
print >>f, "PyObject* PyAST_mod2obj(mod_ty t);"
1019-
print >>f, "mod_ty PyAST_obj2mod(PyObject* ast, PyArena* arena);"
1026+
print >>f, "mod_ty PyAST_obj2mod(PyObject* ast, PyArena* arena, int mode);"
10201027
print >>f, "int PyAST_Check(PyObject* obj);"
10211028
f.close()
10221029

‎Python/Python-ast.c‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5944,13 +5944,20 @@ PyObject* PyAST_mod2obj(mod_ty t)
59445944
return ast2obj_mod(t);
59455945
}
59465946

5947-
mod_ty PyAST_obj2mod(PyObject* ast, PyArena* arena)
5947+
/* mode is 0 for "exec", 1 for "eval" and 2 for "single" input */
5948+
mod_ty PyAST_obj2mod(PyObject* ast, PyArena* arena, int mode)
59485949
{
59495950
mod_ty res;
5951+
PyObject *req_type[] = {(PyObject*)Module_type, (PyObject*)Expression_type,
5952+
(PyObject*)Interactive_type};
5953+
char *req_name[] = {"Module", "Expression", "Interactive"};
5954+
assert(0 <= mode && mode <= 2);
5955+
59505956
init_types();
5951-
if (!PyObject_IsInstance(ast, (PyObject*)mod_type)) {
5952-
PyErr_SetString(PyExc_TypeError, "expected either Module, Interactive "
5953-
"or Expression node");
5957+
5958+
if (!PyObject_IsInstance(ast, req_type[mode])) {
5959+
PyErr_Format(PyExc_TypeError, "expected %s node, got %.400s",
5960+
req_name[mode], Py_TYPE(ast)->tp_name);
59545961
return NULL;
59555962
}
59565963
if (obj2ast_mod(ast, &res, arena) != 0)

‎Python/bltinmodule.c‎

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -466,14 +466,15 @@ builtin_compile(PyObject *self, PyObject *args, PyObject *kwds)
466466
char *str;
467467
char *filename;
468468
char *startstr;
469-
int start;
469+
int mode = -1;
470470
int dont_inherit = 0;
471471
int supplied_flags = 0;
472472
PyCompilerFlags cf;
473473
PyObject *result = NULL, *cmd, *tmp = NULL;
474474
Py_ssize_t length;
475475
static char *kwlist[] = {"source", "filename", "mode", "flags",
476476
"dont_inherit", NULL};
477+
int start[] = {Py_file_input, Py_eval_input, Py_single_input};
477478

478479
if (!PyArg_ParseTupleAndKeywords(args, kwds, "Oss|ii:compile",
479480
kwlist, &cmd, &filename, &startstr,
@@ -495,6 +496,18 @@ builtin_compile(PyObject *self, PyObject *args, PyObject *kwds)
495496
PyEval_MergeCompilerFlags(&cf);
496497
}
497498

499+
if (strcmp(startstr, "exec") == 0)
500+
mode = 0;
501+
else if (strcmp(startstr, "eval") == 0)
502+
mode = 1;
503+
else if (strcmp(startstr, "single") == 0)
504+
mode = 2;
505+
else {
506+
PyErr_SetString(PyExc_ValueError,
507+
"compile() arg 3 must be 'exec', 'eval' or 'single'");
508+
return NULL;
509+
}
510+
498511
if (PyAST_Check(cmd)) {
499512
if (supplied_flags & PyCF_ONLY_AST) {
500513
Py_INCREF(cmd);
@@ -505,7 +518,7 @@ builtin_compile(PyObject *self, PyObject *args, PyObject *kwds)
505518
mod_ty mod;
506519

507520
arena = PyArena_New();
508-
mod = PyAST_obj2mod(cmd, arena);
521+
mod = PyAST_obj2mod(cmd, arena, mode);
509522
if (mod == NULL) {
510523
PyArena_Free(arena);
511524
return NULL;
@@ -526,19 +539,6 @@ builtin_compile(PyObject *self, PyObject *args, PyObject *kwds)
526539
cf.cf_flags |= PyCF_SOURCE_IS_UTF8;
527540
}
528541
#endif
529-
/* XXX: is it possible to pass start to the PyAST_ branch? */
530-
if (strcmp(startstr, "exec") == 0)
531-
start = Py_file_input;
532-
else if (strcmp(startstr, "eval") == 0)
533-
start = Py_eval_input;
534-
else if (strcmp(startstr, "single") == 0)
535-
start = Py_single_input;
536-
else {
537-
PyErr_SetString(PyExc_ValueError,
538-
"compile() arg 3 must be 'exec'"
539-
"or 'eval' or 'single'");
540-
goto cleanup;
541-
}
542542

543543
if (PyObject_AsReadBuffer(cmd, (const void **)&str, &length))
544544
goto cleanup;
@@ -547,7 +547,7 @@ builtin_compile(PyObject *self, PyObject *args, PyObject *kwds)
547547
"compile() expected string without null bytes");
548548
goto cleanup;
549549
}
550-
result = Py_CompileStringFlags(str, filename, start, &cf);
550+
result = Py_CompileStringFlags(str, filename, start[mode], &cf);
551551
cleanup:
552552
Py_XDECREF(tmp);
553553
return result;

0 commit comments

Comments
 (0)