fix codegen: modules and import-from statements

1 files changed, 12 insertions, 28 deletions

diff --git a/util/codegen.py b/util/codegen.py
index bd5d679..b53e225 100644
--- a/util/codegen.py
+++ b/util/codegen.py
@@ -1,27 +1,4 @@
 # -*- coding: utf-8 -*-
-# Copyright (c) 2008, Armin Ronacher
-# All rights reserved.
-# 
-# Redistribution and use in source and binary forms, with or without modification,
-# are permitted provided that the following conditions are met:
-# 
-#   Redistributions of source code must retain the above copyright notice, this
-#   list of conditions and the following disclaimer.
-# 
-#   Redistributions in binary form must reproduce the above copyright notice, this
-#   list of conditions and the following disclaimer in the documentation and/or
-#   other materials provided with the distribution.
-# 
-# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
-# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
-# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
-# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
-# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
-# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
-# ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
-# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
-# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 """
     codegen
     ~~~~~~~
@@ -84,12 +61,14 @@ class SourceGenerator(NodeVisitor):
             self.write('# line: %s' % node.lineno)
             self.new_lines = 1
 
-    def body(self, statements):
-        self.new_line = True
-        self.indentation += 1
+    def body(self, statements, do_indent=True):
+        if do_indent:
+            self.indentation += 1
         for stmt in statements:
+            self.newline()
             self.visit(stmt)
-        self.indentation -= 1
+        if do_indent:
+            self.indentation -= 1
 
     def body_or_else(self, node):
         self.body(node.body)
@@ -149,7 +128,7 @@ class SourceGenerator(NodeVisitor):
         for idx, item in enumerate(node.names):
             if idx:
                 self.write(', ')
-            self.write(item)
+            self.visit(item)
 
     def visit_Import(self, node):
         self.newline(node)
@@ -161,6 +140,9 @@ class SourceGenerator(NodeVisitor):
         self.newline(node)
         self.generic_visit(node)
 
+    def visit_Module(self, node):
+        self.body(node.body, do_indent=False)
+
     def visit_FunctionDef(self, node):
         self.newline(extra=1)
         self.decorators(node)
@@ -201,6 +183,7 @@ class SourceGenerator(NodeVisitor):
                 paren_or_comma()
                 self.write('**')
                 self.visit(node.kwargs)
+        # TODO wtf???
         self.write(have_args and '):' or ':')
         self.body(node.body)
 
@@ -387,6 +370,7 @@ class SourceGenerator(NodeVisitor):
             if idx:
                 self.write(', ')
             self.visit(item)
+        # TODO wtf???
         self.write(idx and ')' or ',)')
 
     def sequence_visit(left, right):