Lib/test/test_ast.py

   1 import sys, itertools, unittest
   2 from test import test_support
   3 import ast
   4
   5 def to_tuple(t):
   6     if t is None or isinstance(t, (basestring, int, long, complex)):
   7         return t
   8     elif isinstance(t, list):
   9         return [to_tuple(e) for e in t]
  10     result = [t.__class__.__name__]
  11     if hasattr(t, 'lineno') and hasattr(t, 'col_offset'):
  12         result.append((t.lineno, t.col_offset))
  13     if t._fields is None:
  14         return tuple(result)
  15     for f in t._fields:
  16         result.append(to_tuple(getattr(t, f)))
  17     return tuple(result)
  18
  19
  20 # These tests are compiled through "exec"
  21 # There should be atleast one test per statement
  22 exec_tests = [
  23     # FunctionDef
  24     "def f(): pass",
  25     # ClassDef
  26     "class C:pass",
  27     # Return
  28     "def f():return 1",
  29     # Delete
  30     "del v",
  31     # Assign
  32     "v = 1",
  33     # AugAssign
  34     "v += 1",
  35     # Print
  36     "print >>f, 1, ",
  37     # For
  38     "for v in v:pass",
  39     # While
  40     "while v:pass",
  41     # If
  42     "if v:pass",
  43     # Raise
  44     "raise Exception, 'string'",
  45     # TryExcept
  46     "try:\n  pass\nexcept Exception:\n  pass",
  47     # TryFinally
  48     "try:\n  pass\nfinally:\n  pass",
  49     # Assert
  50     "assert v",
  51     # Import
  52     "import sys",
  53     # ImportFrom
  54     "from sys import v",
  55     # Exec
  56     "exec 'v'",
  57     # Global
  58     "global v",
  59     # Expr
  60     "1",
  61     # Pass,
  62     "pass",
  63     # Break
  64     "break",
  65     # Continue
  66     "continue",
  67     # for statements with naked tuples (see http://bugs.python.org/issue6704)
  68     "for a,b in c: pass",
  69     "[(a,b) for a,b in c]",
  70     "((a,b) for a,b in c)",
  71 ]
  72
  73 # These are compiled through "single"
  74 # because of overlap with "eval", it just tests what
  75 # can't be tested with "eval"
  76 single_tests = [
  77     "1+2"
  78 ]
  79
  80 # These are compiled through "eval"
  81 # It should test all expressions
  82 eval_tests = [
  83   # BoolOp
  84   "a and b",
  85   # BinOp
  86   "a + b",
  87   # UnaryOp
  88   "not v",
  89   # Lambda
  90   "lambda:None",
  91   # Dict
  92   "{ 1:2 }",
  93   # ListComp
  94   "[a for b in c if d]",
  95   # GeneratorExp
  96   "(a for b in c if d)",
  97   # Yield - yield expressions can't work outside a function
  98   #
  99   # Compare
 100   "1 < 2 < 3",
 101   # Call
 102   "f(1,2,c=3,*d,**e)",
 103   # Repr
 104   "`v`",
 105   # Num
 106   "10L",
 107   # Str
 108   "'string'",
 109   # Attribute
 110   "a.b",
 111   # Subscript
 112   "a[b:c]",
 113   # Name
 114   "v",
 115   # List
 116   "[1,2,3]",
 117   # Tuple
 118   "1,2,3",
 119   # Combination
 120   "a.b.c.d(a.b[1:2])",
 121
 122 ]
 123
 124 # TODO: expr_context, slice, boolop, operator, unaryop, cmpop, comprehension
 125 # excepthandler, arguments, keywords, alias
 126
 127 class AST_Tests(unittest.TestCase):
 128
 129     def _assertTrueorder(self, ast_node, parent_pos):
 130         if not isinstance(ast_node, ast.AST) or ast_node._fields is None:
 131             return
 132         if isinstance(ast_node, (ast.expr, ast.stmt, ast.excepthandler)):
 133             node_pos = (ast_node.lineno, ast_node.col_offset)
 134             self.assertTrue(node_pos >= parent_pos)
 135             parent_pos = (ast_node.lineno, ast_node.col_offset)
 136         for name in ast_node._fields:
 137             value = getattr(ast_node, name)
 138             if isinstance(value, list):
 139                 for child in value:
 140                     self._assertTrueorder(child, parent_pos)
 141             elif value is not None:
 142                 self._assertTrueorder(value, parent_pos)
 143
 144     def test_snippets(self):
 145         for input, output, kind in ((exec_tests, exec_results, "exec"),
 146                                     (single_tests, single_results, "single"),
 147                                     (eval_tests, eval_results, "eval")):
 148             for i, o in itertools.izip(input, output):
 149                 ast_tree = compile(i, "?", kind, ast.PyCF_ONLY_AST)
 150                 self.assertEquals(to_tuple(ast_tree), o)
 151                 self._assertTrueorder(ast_tree, (0, 0))
 152
 153     def test_slice(self):
 154         slc = ast.parse("x[::]").body[0].value.slice
 155         self.assertIsNone(slc.upper)
 156         self.assertIsNone(slc.lower)
 157         self.assertIsInstance(slc.step, ast.Name)
 158         self.assertEqual(slc.step.id, "None")
 159
 160     def test_from_import(self):
 161         im = ast.parse("from . import y").body[0]
 162         self.assertIsNone(im.module)
 163
 164     def test_base_classes(self):
 165         self.assertTrue(issubclass(ast.For, ast.stmt))
 166         self.assertTrue(issubclass(ast.Name, ast.expr))
 167         self.assertTrue(issubclass(ast.stmt, ast.AST))
 168         self.assertTrue(issubclass(ast.expr, ast.AST))
 169         self.assertTrue(issubclass(ast.comprehension, ast.AST))
 170         self.assertTrue(issubclass(ast.Gt, ast.AST))
 171
 172     def test_nodeclasses(self):
 173         x = ast.BinOp(1, 2, 3, lineno=0)
 174         self.assertEquals(x.left, 1)
 175         self.assertEquals(x.op, 2)
 176         self.assertEquals(x.right, 3)
 177         self.assertEquals(x.lineno, 0)
 178
 179         # node raises exception when not given enough arguments
 180         self.assertRaises(TypeError, ast.BinOp, 1, 2)
 181
 182         # can set attributes through kwargs too
 183         x = ast.BinOp(left=1, op=2, right=3, lineno=0)
 184         self.assertEquals(x.left, 1)
 185         self.assertEquals(x.op, 2)
 186         self.assertEquals(x.right, 3)
 187         self.assertEquals(x.lineno, 0)
 188
 189         # this used to fail because Sub._fields was None
 190         x = ast.Sub()
 191
 192     def test_pickling(self):
 193         import pickle
 194         mods = [pickle]
 195         try:
 196             import cPickle
 197             mods.append(cPickle)
 198         except ImportError:
 199             pass
 200         protocols = [0, 1, 2]
 201         for mod in mods:
 202             for protocol in protocols:
 203                 for ast in (compile(i, "?", "exec", 0x400) for i in exec_tests):
 204                     ast2 = mod.loads(mod.dumps(ast, protocol))
 205                     self.assertEquals(to_tuple(ast2), to_tuple(ast))
 206
 207
 208 class ASTHelpers_Test(unittest.TestCase):
 209
 210     def test_parse(self):
 211         a = ast.parse('foo(1 + 1)')
 212         b = compile('foo(1 + 1)', '<unknown>', 'exec', ast.PyCF_ONLY_AST)
 213         self.assertEqual(ast.dump(a), ast.dump(b))
 214
 215     def test_dump(self):
 216         node = ast.parse('spam(eggs, "and cheese")')
 217         self.assertEqual(ast.dump(node),
 218             "Module(body=[Expr(value=Call(func=Name(id='spam', ctx=Load()), "
 219             "args=[Name(id='eggs', ctx=Load()), Str(s='and cheese')], "
 220             "keywords=[], starargs=None, kwargs=None))])"
 221         )
 222         self.assertEqual(ast.dump(node, annotate_fields=False),
 223             "Module([Expr(Call(Name('spam', Load()), [Name('eggs', Load()), "
 224             "Str('and cheese')], [], None, None))])"
 225         )
 226         self.assertEqual(ast.dump(node, include_attributes=True),
 227             "Module(body=[Expr(value=Call(func=Name(id='spam', ctx=Load(), "
 228             "lineno=1, col_offset=0), args=[Name(id='eggs', ctx=Load(), "
 229             "lineno=1, col_offset=5), Str(s='and cheese', lineno=1, "
 230             "col_offset=11)], keywords=[], starargs=None, kwargs=None, "
 231             "lineno=1, col_offset=0), lineno=1, col_offset=0)])"
 232         )
 233
 234     def test_copy_location(self):
 235         src = ast.parse('1 + 1', mode='eval')
 236         src.body.right = ast.copy_location(ast.Num(2), src.body.right)
 237         self.assertEqual(ast.dump(src, include_attributes=True),
 238             'Expression(body=BinOp(left=Num(n=1, lineno=1, col_offset=0), '
 239             'op=Add(), right=Num(n=2, lineno=1, col_offset=4), lineno=1, '
 240             'col_offset=0))'
 241         )
 242
 243     def test_fix_missing_locations(self):
 244         src = ast.parse('write("spam")')
 245         src.body.append(ast.Expr(ast.Call(ast.Name('spam', ast.Load()),
 246                                           [ast.Str('eggs')], [], None, None)))
 247         self.assertEqual(src, ast.fix_missing_locations(src))
 248         self.assertEqual(ast.dump(src, include_attributes=True),
 249             "Module(body=[Expr(value=Call(func=Name(id='write', ctx=Load(), "
 250             "lineno=1, col_offset=0), args=[Str(s='spam', lineno=1, "
 251             "col_offset=6)], keywords=[], starargs=None, kwargs=None, "
 252             "lineno=1, col_offset=0), lineno=1, col_offset=0), "
 253             "Expr(value=Call(func=Name(id='spam', ctx=Load(), lineno=1, "
 254             "col_offset=0), args=[Str(s='eggs', lineno=1, col_offset=0)], "
 255             "keywords=[], starargs=None, kwargs=None, lineno=1, "
 256             "col_offset=0), lineno=1, col_offset=0)])"
 257         )
 258
 259     def test_increment_lineno(self):
 260         src = ast.parse('1 + 1', mode='eval')
 261         self.assertEqual(ast.increment_lineno(src, n=3), src)
 262         self.assertEqual(ast.dump(src, include_attributes=True),
 263             'Expression(body=BinOp(left=Num(n=1, lineno=4, col_offset=0), '
 264             'op=Add(), right=Num(n=1, lineno=4, col_offset=4), lineno=4, '
 265             'col_offset=0))'
 266         )
 267
 268     def test_iter_fields(self):
 269         node = ast.parse('foo()', mode='eval')
 270         d = dict(ast.iter_fields(node.body))
 271         self.assertEqual(d.pop('func').id, 'foo')
 272         self.assertEqual(d, {'keywords': [], 'kwargs': None,
 273                              'args': [], 'starargs': None})
 274
 275     def test_iter_child_nodes(self):
 276         node = ast.parse("spam(23, 42, eggs='leek')", mode='eval')
 277         self.assertEqual(len(list(ast.iter_child_nodes(node.body))), 4)
 278         iterator = ast.iter_child_nodes(node.body)
 279         self.assertEqual(next(iterator).id, 'spam')
 280         self.assertEqual(next(iterator).n, 23)
 281         self.assertEqual(next(iterator).n, 42)
 282         self.assertEqual(ast.dump(next(iterator)),
 283             "keyword(arg='eggs', value=Str(s='leek'))"
 284         )
 285
 286     def test_get_docstring(self):
 287         node = ast.parse('def foo():\n  """line one\n  line two"""')
 288         self.assertEqual(ast.get_docstring(node.body[0]),
 289                          'line one\nline two')
 290
 291     def test_literal_eval(self):
 292         self.assertEqual(ast.literal_eval('[1, 2, 3]'), [1, 2, 3])
 293         self.assertEqual(ast.literal_eval('{"foo": 42}'), {"foo": 42})
 294         self.assertEqual(ast.literal_eval('(True, False, None)'), (True, False, None))
 295         self.assertRaises(ValueError, ast.literal_eval, 'foo()')
 296
 297     def test_literal_eval_issue4907(self):
 298         self.assertEqual(ast.literal_eval('2j'), 2j)
 299         self.assertEqual(ast.literal_eval('10 + 2j'), 10 + 2j)
 300         self.assertEqual(ast.literal_eval('1.5 - 2j'), 1.5 - 2j)
 301         self.assertRaises(ValueError, ast.literal_eval, '2 + (3 + 4j)')
 302
 303
 304 def test_main():
 305     with test_support.check_py3k_warnings(("backquote not supported",
 306                                              SyntaxWarning)):
 307         test_support.run_unittest(AST_Tests, ASTHelpers_Test)
 308
 309 def main():
 310     if __name__ != '__main__':
 311         return
 312     if sys.argv[1:] == ['-g']:
 313         for statements, kind in ((exec_tests, "exec"), (single_tests, "single"),
 314                                  (eval_tests, "eval")):
 315             print kind+"_results = ["
 316             for s in statements:
 317                 print repr(to_tuple(compile(s, "?", kind, 0x400)))+","
 318             print "]"
 319         print "main()"
 320         raise SystemExit
 321     test_main()
 322
 323 #### EVERYTHING BELOW IS GENERATED #####
 324 exec_results = [
 325 ('Module', [('FunctionDef', (1, 0), 'f', ('arguments', [], None, None, []), [('Pass', (1, 9))], [])]),
 326 ('Module', [('ClassDef', (1, 0), 'C', [], [('Pass', (1, 8))], [])]),
 327 ('Module', [('FunctionDef', (1, 0), 'f', ('arguments', [], None, None, []), [('Return', (1, 8), ('Num', (1, 15), 1))], [])]),
 328 ('Module', [('Delete', (1, 0), [('Name', (1, 4), 'v', ('Del',))])]),
 329 ('Module', [('Assign', (1, 0), [('Name', (1, 0), 'v', ('Store',))], ('Num', (1, 4), 1))]),
 330 ('Module', [('AugAssign', (1, 0), ('Name', (1, 0), 'v', ('Store',)), ('Add',), ('Num', (1, 5), 1))]),
 331 ('Module', [('Print', (1, 0), ('Name', (1, 8), 'f', ('Load',)), [('Num', (1, 11), 1)], False)]),
 332 ('Module', [('For', (1, 0), ('Name', (1, 4), 'v', ('Store',)), ('Name', (1, 9), 'v', ('Load',)), [('Pass', (1, 11))], [])]),
 333 ('Module', [('While', (1, 0), ('Name', (1, 6), 'v', ('Load',)), [('Pass', (1, 8))], [])]),
 334 ('Module', [('If', (1, 0), ('Name', (1, 3), 'v', ('Load',)), [('Pass', (1, 5))], [])]),
 335 ('Module', [('Raise', (1, 0), ('Name', (1, 6), 'Exception', ('Load',)), ('Str', (1, 17), 'string'), None)]),
 336 ('Module', [('TryExcept', (1, 0), [('Pass', (2, 2))], [('ExceptHandler', (3, 0), ('Name', (3, 7), 'Exception', ('Load',)), None, [('Pass', (4, 2))])], [])]),
 337 ('Module', [('TryFinally', (1, 0), [('Pass', (2, 2))], [('Pass', (4, 2))])]),
 338 ('Module', [('Assert', (1, 0), ('Name', (1, 7), 'v', ('Load',)), None)]),
 339 ('Module', [('Import', (1, 0), [('alias', 'sys', None)])]),
 340 ('Module', [('ImportFrom', (1, 0), 'sys', [('alias', 'v', None)], 0)]),
 341 ('Module', [('Exec', (1, 0), ('Str', (1, 5), 'v'), None, None)]),
 342 ('Module', [('Global', (1, 0), ['v'])]),
 343 ('Module', [('Expr', (1, 0), ('Num', (1, 0), 1))]),
 344 ('Module', [('Pass', (1, 0))]),
 345 ('Module', [('Break', (1, 0))]),
 346 ('Module', [('Continue', (1, 0))]),
 347 ('Module', [('For', (1, 0), ('Tuple', (1, 4), [('Name', (1, 4), 'a', ('Store',)), ('Name', (1, 6), 'b', ('Store',))], ('Store',)), ('Name', (1, 11), 'c', ('Load',)), [('Pass', (1, 14))], [])]),
 348 ('Module', [('Expr', (1, 0), ('ListComp', (1, 1), ('Tuple', (1, 2), [('Name', (1, 2), 'a', ('Load',)), ('Name', (1, 4), 'b', ('Load',))], ('Load',)), [('comprehension', ('Tuple', (1, 11), [('Name', (1, 11), 'a', ('Store',)), ('Name', (1, 13), 'b', ('Store',))], ('Store',)), ('Name', (1, 18), 'c', ('Load',)), [])]))]),
 349 ('Module', [('Expr', (1, 0), ('GeneratorExp', (1, 1), ('Tuple', (1, 2), [('Name', (1, 2), 'a', ('Load',)), ('Name', (1, 4), 'b', ('Load',))], ('Load',)), [('comprehension', ('Tuple', (1, 11), [('Name', (1, 11), 'a', ('Store',)), ('Name', (1, 13), 'b', ('Store',))], ('Store',)), ('Name', (1, 18), 'c', ('Load',)), [])]))]),
 350 ]
 351 single_results = [
 352 ('Interactive', [('Expr', (1, 0), ('BinOp', (1, 0), ('Num', (1, 0), 1), ('Add',), ('Num', (1, 2), 2)))]),
 353 ]
 354 eval_results = [
 355 ('Expression', ('BoolOp', (1, 0), ('And',), [('Name', (1, 0), 'a', ('Load',)), ('Name', (1, 6), 'b', ('Load',))])),
 356 ('Expression', ('BinOp', (1, 0), ('Name', (1, 0), 'a', ('Load',)), ('Add',), ('Name', (1, 4), 'b', ('Load',)))),
 357 ('Expression', ('UnaryOp', (1, 0), ('Not',), ('Name', (1, 4), 'v', ('Load',)))),
 358 ('Expression', ('Lambda', (1, 0), ('arguments', [], None, None, []), ('Name', (1, 7), 'None', ('Load',)))),
 359 ('Expression', ('Dict', (1, 0), [('Num', (1, 2), 1)], [('Num', (1, 4), 2)])),
 360 ('Expression', ('ListComp', (1, 1), ('Name', (1, 1), 'a', ('Load',)), [('comprehension', ('Name', (1, 7), 'b', ('Store',)), ('Name', (1, 12), 'c', ('Load',)), [('Name', (1, 17), 'd', ('Load',))])])),
 361 ('Expression', ('GeneratorExp', (1, 1), ('Name', (1, 1), 'a', ('Load',)), [('comprehension', ('Name', (1, 7), 'b', ('Store',)), ('Name', (1, 12), 'c', ('Load',)), [('Name', (1, 17), 'd', ('Load',))])])),
 362 ('Expression', ('Compare', (1, 0), ('Num', (1, 0), 1), [('Lt',), ('Lt',)], [('Num', (1, 4), 2), ('Num', (1, 8), 3)])),
 363 ('Expression', ('Call', (1, 0), ('Name', (1, 0), 'f', ('Load',)), [('Num', (1, 2), 1), ('Num', (1, 4), 2)], [('keyword', 'c', ('Num', (1, 8), 3))], ('Name', (1, 11), 'd', ('Load',)), ('Name', (1, 15), 'e', ('Load',)))),
 364 ('Expression', ('Repr', (1, 0), ('Name', (1, 1), 'v', ('Load',)))),
 365 ('Expression', ('Num', (1, 0), 10L)),
 366 ('Expression', ('Str', (1, 0), 'string')),
 367 ('Expression', ('Attribute', (1, 0), ('Name', (1, 0), 'a', ('Load',)), 'b', ('Load',))),
 368 ('Expression', ('Subscript', (1, 0), ('Name', (1, 0), 'a', ('Load',)), ('Slice', ('Name', (1, 2), 'b', ('Load',)), ('Name', (1, 4), 'c', ('Load',)), None), ('Load',))),
 369 ('Expression', ('Name', (1, 0), 'v', ('Load',))),
 370 ('Expression', ('List', (1, 0), [('Num', (1, 1), 1), ('Num', (1, 3), 2), ('Num', (1, 5), 3)], ('Load',))),
 371 ('Expression', ('Tuple', (1, 0), [('Num', (1, 0), 1), ('Num', (1, 2), 2), ('Num', (1, 4), 3)], ('Load',))),
 372 ('Expression', ('Call', (1, 0), ('Attribute', (1, 0), ('Attribute', (1, 0), ('Attribute', (1, 0), ('Name', (1, 0), 'a', ('Load',)), 'b', ('Load',)), 'c', ('Load',)), 'd', ('Load',)), [('Subscript', (1, 8), ('Attribute', (1, 8), ('Name', (1, 8), 'a', ('Load',)), 'b', ('Load',)), ('Slice', ('Num', (1, 12), 1), ('Num', (1, 14), 2), None), ('Load',))], [], None, None)),
 373 ]
 374 main()