make the patched compile() work with AST

--HG--
branch : trunk

py/code/source.py

@@ -4,6 +4,14 @@ import inspect, tokenize
 import py
 cpy_compile = compile 
 
+try:
+    import _ast
+    from _ast import PyCF_ONLY_AST as _AST_FLAG
+except ImportError:
+    _AST_FLAG = 0
+    _ast = None
+
+
 class Source(object):
     """ a mutable object holding a source code fragment,
         possibly deindenting it.
@@ -195,6 +203,8 @@ class Source(object):
             newex.text = ex.text
             raise newex
         else:
+            if flag & _AST_FLAG:
+                return co
             co_filename = MyStr(filename)
             co_filename.__source__ = self
             return py.code.Code(co).new(rec=1, co_filename=co_filename) 
@@ -213,6 +223,9 @@ def compile_(source, filename=None, mode='exec', flags=
         also have this special subclass-of-string
         filename.
     """
+    if _ast is not None and isinstance(source, _ast.AST):
+        # XXX should Source support having AST?
+        return cpy_compile(source, filename, mode, flags, dont_inherit)
     _genframe = sys._getframe(1) # the caller
     s = Source(source)
     co = s.compile(filename, mode, flags, _genframe=_genframe)

py/code/testing/test_source.py

@@ -191,6 +191,15 @@ class TestSourceParsingAndCompiling:
         assert len(source) == 9
         assert source.getstatementrange(5) == (0, 9)
 
+    def test_compile_to_ast(self):
+        if sys.version_info < (2, 5):
+            py.test.skip("requires Python 2.5")
+        import _ast
+        source = Source("x = 4")
+        mod = source.compile(flag=_ast.PyCF_ONLY_AST)
+        assert isinstance(mod, _ast.Module)
+        compile(mod, "<filename>", "exec")
+
     def test_compile_and_getsource(self):
         co = self.source.compile()
         py.builtin.exec_(co, globals())