diff --git a/codegen.py b/codegen.py index 113d9be..07ac543 100644 --- a/codegen.py +++ b/codegen.py @@ -169,7 +169,10 @@ def visit_ImportFrom(self, node): for idx, item in enumerate(node.names): if idx: self.write(', ') - self.write(item) + self.write(item.name) + if item.asname is not None: + self.write(' as ') + self.write(item.asname) def visit_Import(self, node): self.newline(node) @@ -297,7 +300,7 @@ def visit_Print(self, node): def visit_Delete(self, node): self.newline(node) self.write('del ') - for idx, target in enumerate(node): + for idx, target in enumerate(node.targets): if idx: self.write(', ') self.visit(target) @@ -309,6 +312,16 @@ def visit_TryExcept(self, node): for handler in node.handlers: self.visit(handler) + def visit_ExceptHandler(self, node): + self.newline(node) + self.write('except ') + self.visit(node.type) + if node.name is not None: + self.write(', ') + self.visit(node.name) + self.write(':') + self.body(node.body) + def visit_TryFinally(self, node): self.newline(node) self.write('try:') diff --git a/test/test_codegen.py b/test/test_codegen.py new file mode 100644 index 0000000..91de0c0 --- /dev/null +++ b/test/test_codegen.py @@ -0,0 +1,30 @@ +import codegen +import ast + +def to_ast_and_back_again(source): + return codegen.to_source(ast.parse(source)) + +def test_del(): + source = "del l[0]" + assert source == to_ast_and_back_again(source) + source = "del obj.x" + assert source == to_ast_and_back_again(source) + +def test_try_expect(): + source = ("try:\n" + " '#'[2]\n" + "except IndexError:\n" + " print 'What did you expect?!'") + assert source == to_ast_and_back_again(source) + source = ("try:\n" + " l = []\n" + " l[1]\n" + "except IndexError, index_error:\n" + " print index_error") + assert source == to_ast_and_back_again(source) + +def test_import(): + source = "import intertools as iterators" + assert source == to_ast_and_back_again(source) + source = "from math import floor as fl, ceil as cl" + assert source == to_ast_and_back_again(source)