codegen dev...

1 files changed, 39 insertions, 6 deletions

diff --git a/test/ast.py b/test/ast.py
index b5d90c9..3b1ec8b 100644
--- a/test/ast.py
+++ b/test/ast.py
@@ -1,6 +1,6 @@
 import unittest
 
-from mako import ast, util
+from mako import ast, util, exceptions
 from compiler import parse
 
 class AstParseTest(unittest.TestCase):
@@ -23,31 +23,64 @@ for lar in (1,2,3):
 print "hello world, ", a, b
 print "Another expr", c
 """
-        parsed = ast.PythonCode(code)
+        parsed = ast.PythonCode(code, 0, 0)
         assert parsed.declared_identifiers == util.Set(['a','b','c', 'g', 'h', 'i', 'u', 'k', 'j', 'gh', 'lar'])
         assert parsed.undeclared_identifiers == util.Set(['x', 'q', 'foo', 'gah', 'blah'])
     
-        parsed = ast.PythonCode("x + 5 * (y-z)")
+        parsed = ast.PythonCode("x + 5 * (y-z)", 0, 0)
         assert parsed.undeclared_identifiers == util.Set(['x', 'y', 'z'])
         assert parsed.declared_identifiers == util.Set()
 
     def test_locate_identifiers_2(self):
         code = """
+import foobar
+from lala import hoho, yaya
+import bleep as foo
 result = []
 data = get_data()
 for x in data:
     result.append(x+7)
 """
-        parsed = ast.PythonCode(code)
+        parsed = ast.PythonCode(code, 0, 0)
+        print parsed.declared_identifiers
         assert parsed.undeclared_identifiers == util.Set(['get_data'])
-        assert parsed.declared_identifiers == util.Set(['result', 'data', 'x'])
+        assert parsed.declared_identifiers == util.Set(['result', 'data', 'x', 'hoho', 'foobar', 'foo', 'yaya'])
+
+    def test_no_global_imports(self):
+        code = """
+from foo import *
+import x as bar
+"""
+        try:
+            parsed = ast.PythonCode(code, 0, 0)
+            assert False
+        except exceptions.CompileException, e:
+            assert str(e).startswith("'import *' is not supported")
+            
+    def test_python_fragment(self):
+        parsed = ast.PythonFragment("for x in foo:", 0, 0)
+        assert parsed.declared_identifiers == util.Set(['x'])
+        assert parsed.undeclared_identifiers == util.Set(['foo'])
+        
+        parsed = ast.PythonFragment("try:", 0, 0)
+
+        parsed = ast.PythonFragment("except (MyException, e):", 0, 0)
+        assert parsed.declared_identifiers == util.Set(['e'])
+        assert parsed.undeclared_identifiers == util.Set()
         
     def test_function_decl(self):
         """test getting the arguments from a function"""
         code = "def foo(a, b, c=None, d='hi', e=x, f=y+7):pass"
-        parsed = ast.FunctionDecl(code)
+        parsed = ast.FunctionDecl(code, 0, 0)
         assert parsed.funcname=='foo'
         assert parsed.argnames==['a', 'b', 'c', 'd', 'e', 'f']
+
+    def test_function_decl_2(self):
+        """test getting the arguments from a function"""
+        code = "def foo(a, b, c=None, *args, **kwargs):pass"
+        parsed = ast.FunctionDecl(code, 0, 0)
+        assert parsed.funcname=='foo'
+        assert parsed.argnames==['a', 'b', 'c', 'args', 'kwargs']
     
     def test_expr_generate(self):
         """test the round trip of expressions to AST back to python source"""