[pypy-svn] r78323 - in pypy/branch/fast-forward/pypy: interpreter module/__builtin__/test

[afa at codespeak.net]

Author: afa
Date: Tue Oct 26 23:03:05 2010
New Revision: 78323

Modified:
   pypy/branch/fast-forward/pypy/interpreter/pycompiler.py
   pypy/branch/fast-forward/pypy/module/__builtin__/test/test_builtin.py
Log:
Check node type when compiling an AST tree


Modified: pypy/branch/fast-forward/pypy/interpreter/pycompiler.py
==============================================================================
--- pypy/branch/fast-forward/pypy/interpreter/pycompiler.py	(original)
+++ pypy/branch/fast-forward/pypy/interpreter/pycompiler.py	Tue Oct 26 23:03:05 2010
@@ -6,7 +6,7 @@
 from pypy.interpreter import pycode
 from pypy.interpreter.pyparser import future, pyparse, error as parseerror
 from pypy.interpreter.astcompiler import (astbuilder, codegen, consts, misc,
-                                          optimize)
+                                          optimize, ast)
 from pypy.interpreter.error import OperationError
 
 
@@ -107,6 +107,18 @@
         self.compiler_flags = self.future_flags.allowed_flags
 
     def compile_ast(self, node, filename, mode, flags):
+        if mode == 'eval':
+            check = isinstance(node, ast.Expression)
+        elif mode == 'exec':
+            check = isinstance(node, ast.Module)
+        elif mode == 'input':
+            check = isinstance(node, ast.Interactive)
+        else:
+            check = True
+        if not check:
+            raise OperationError(self.space.w_TypeError, self.space.wrap(
+                "invalid node type"))
+
         future_pos = misc.parse_future(node)
         info = pyparse.CompileInfo(filename, mode, flags, future_pos)
         return self._compile_ast(node, info)

Modified: pypy/branch/fast-forward/pypy/module/__builtin__/test/test_builtin.py
==============================================================================
--- pypy/branch/fast-forward/pypy/module/__builtin__/test/test_builtin.py	(original)
+++ pypy/branch/fast-forward/pypy/module/__builtin__/test/test_builtin.py	Tue Oct 26 23:03:05 2010
@@ -453,7 +453,15 @@
     def test_unicode_encoding_compile(self):
         code = u"# -*- coding: utf-8 -*-\npass\n"
         raises(SyntaxError, compile, code, "tmp", "exec")
-            
+
+    def test_recompile_ast(self):
+        import _ast
+        # raise exception when node type doesn't match with compile mode
+        co1 = compile('print 1', '<string>', 'exec', _ast.PyCF_ONLY_AST)
+        raises(TypeError, compile, co1, '<ast>', 'eval')
+        co2 = compile('1+1', '<string>', 'eval', _ast.PyCF_ONLY_AST)
+        compile(co2, '<ast>', 'eval')
+
     def test_isinstance(self):
         assert isinstance(5, int)
         assert isinstance(5, object)