🐛 Fix suite transformer

Make sure to visit non-suite child nodes. Added a consistent way to
add child nodes.

Author: dflook

src/python_minifier/module_printer.py

@@ -494,7 +494,7 @@ def visit_FunctionDef(self, node, is_async=False):
         self.visit_arguments(node.args)
         self.code += ')'
 
-        if hasattr(node, 'returns') and node.returns:
+        if hasattr(node, 'returns') and node.returns is not None:
             self.code += '->'
             self._expression(node.returns)
             self.code += ':'

src/python_minifier/transforms/combine_imports.py

@@ -11,7 +11,7 @@ class CombineImports(SuiteTransformer):
 
     """
 
-    def _combine_import(self, node_list):
+    def _combine_import(self, node_list, parent):
 
         alias = []
 
@@ -20,15 +20,17 @@ def _combine_import(self, node_list):
                 alias += statement.names
             else:
                 if alias:
-                    yield ast.Import(names=alias)
+                    yield self.add_child(ast.Import(names=alias),
+                                         parent=parent)
                     alias = []
 
                 yield statement
 
         if alias:
-            yield ast.Import(names=alias)
+            yield self.add_child(ast.Import(names=alias),
+                                 parent=parent)
 
-    def _combine_import_from(self, node_list):
+    def _combine_import_from(self, node_list, parent):
 
         prev_import = None
         alias = []
@@ -54,16 +56,18 @@ def combine(statement):
                 alias += statement.names
             else:
                 if alias:
-                    yield ast.ImportFrom(module=prev_import.module, names=alias, level=prev_import.level)
+                    yield self.add_child(ast.ImportFrom(module=prev_import.module, names=alias, level=prev_import.level),
+                                         parent=parent)
                     alias = []
 
                 yield statement
 
         if alias:
-            yield ast.ImportFrom(module=prev_import.module, names=alias, level=prev_import.level)
+            yield self.add_child(ast.ImportFrom(module=prev_import.module, names=alias, level=prev_import.level),
+                                 parent=parent)
 
     def suite(self, node_list, parent):
-        a = list(self._combine_import(node_list))
-        b = list(self._combine_import_from(a))
+        a = list(self._combine_import(node_list, parent))
+        b = list(self._combine_import_from(a, parent))
 
         return [self.visit(n) for n in b]

src/python_minifier/transforms/remove_annotations.py

@@ -1,7 +1,6 @@
 import ast
 import sys
 
-from python_minifier.rename.mapper import add_parent
 from python_minifier.transforms.suite_transformer import SuiteTransformer
 
 
@@ -16,22 +15,21 @@ def __call__(self, node):
         return self.visit(node)
 
     def visit_FunctionDef(self, node):
+        node.args = self.visit_arguments(node.args)
+        node.body = self.suite(node.body, parent=node)
+        node.decorator_list = [self.visit(d) for d in node.decorator_list]
+
         if hasattr(node, 'returns'):
             node.returns = None
-        node.body = [self.visit(a) for a in node.body]
-
-        if node.args:
-            node.args = self.visit_arguments(node.args)
 
         return node
 
-    def visit_AsyncFunctionDef(self, node):
-        return self.visit_FunctionDef(node)
-
     def visit_arguments(self, node):
+        assert isinstance(node, ast.arguments)
 
         if node.args:
             node.args = [self.visit_arg(a) for a in node.args]
+
         if hasattr(node, 'kwonlyargs') and node.kwonlyargs:
             node.kwonlyargs = [self.visit_arg(a) for a in node.kwonlyargs]
 
@@ -76,15 +74,11 @@ def is_dataclass_field(node):
         if is_dataclass_field(node):
             return node
         elif node.value:
-            assign = ast.Assign([node.target], node.value)
-            assign.parent = node.parent
-            assign.namespace = node.namespace
-            return  assign
+            return self.add_child(ast.Assign([node.target], node.value), parent=node.parent)
         else:
             # Valueless annotations cause the interpreter to treat the variable as a local.
             # I don't know of another way to do that without assigning to it, so
             # keep it as an AnnAssign, but replace the annotation with '0'
 
-            node.annotation = ast.Num(0)
-            add_parent(node, parent=node.parent, namespace=node.namespace)
+            node.annotation = self.add_child(ast.Num(0), parent=node.parent)
             return node

src/python_minifier/transforms/remove_literal_statements.py

@@ -1,6 +1,5 @@
 import ast
 
-from python_minifier.rename.mapper import add_parent
 from python_minifier.transforms.suite_transformer import SuiteTransformer
 
 
@@ -46,24 +45,15 @@ def is_literal_statement(self, node):
         if not isinstance(node, ast.Expr):
             return False
 
-        if (
-            isinstance(node.value, (ast.Num, ast.Str, ast.NameConstant))
-            or node.value.__class__.__name__ == 'Constant'
-            or node.value.__class__.__name__ == 'Bytes'
-        ):
-            return True
-
-        return False
+        return self.is_node(node.value, (ast.Num, ast.Str, 'NameConstant', 'Bytes'))
 
     def suite(self, node_list, parent):
-        without_literals = [self.visit(a) for a in filter(lambda n: not self.is_literal_statement(n), node_list)]
+        without_literals = [self.visit(n) for n in node_list if not self.is_literal_statement(n)]
 
         if len(without_literals) == 0:
             if isinstance(parent, ast.Module):
                 return []
             else:
-                expr = ast.Expr(value=ast.Num(0))
-                add_parent(expr, parent=parent, namespace=parent.namespace)
-                return [expr]
+                return [self.add_child(ast.Expr(value=ast.Num(0)), parent=parent)]
 
         return without_literals

src/python_minifier/transforms/remove_pass.py

@@ -1,7 +1,6 @@
 import ast
 
 from python_minifier.transforms.suite_transformer import SuiteTransformer
-from python_minifier.rename.mapper import add_parent
 
 
 class RemovePass(SuiteTransformer):
@@ -15,14 +14,12 @@ def __call__(self, node):
         return self.visit(node)
 
     def suite(self, node_list, parent):
-        without_pass = [self.visit(a) for a in filter(lambda n: not isinstance(n, ast.Pass), node_list)]
+        without_pass = [self.visit(a) for a in filter(lambda n: not self.is_node(n, ast.Pass), node_list)]
 
         if len(without_pass) == 0:
             if isinstance(parent, ast.Module):
                 return []
             else:
-                expr = ast.Expr(value=ast.Num(0))
-                add_parent(expr, parent=parent, namespace=parent.namespace)
-                return [expr]
+                return [self.add_child(ast.Expr(value=ast.Num(0)), parent=parent)]
 
         return without_pass

src/python_minifier/transforms/suite_transformer.py

@@ -1,5 +1,8 @@
 import ast
 
+from python_minifier.rename.mapper import add_parent
+
+
 class NodeVisitor(object):
 
     def visit(self, node):
@@ -35,6 +38,53 @@ def visit_Constant(self, node):
         visitor = getattr(self, method, self.generic_visit)
         return visitor(node)
 
+    def is_node(self, node, types):
+        """
+        Is a node one of the specified node types
+
+        A node type may be an actual ast class, or a string naming one.
+        types is a single node type or an iterable of many.
+
+        If a node_type specified a specific Constant type (Str, Bytes, Num etc),
+        returns true for Constant nodes of the correct type.
+
+        :type node: ast.AST
+        :param types:
+        :rtype: bool
+        """
+
+
+        if not isinstance(types, tuple):
+            types = types,
+
+        actual_types = []
+        for type in types:
+            if isinstance(type, str):
+                node_type = getattr(ast, type, None)
+                if node_type is not None:
+                    actual_types.append(node_type)
+            else:
+                actual_types.append(type)
+
+        if isinstance(node, tuple(actual_types)):
+            return True
+
+        if hasattr(ast, 'Constant') and isinstance(node, ast.Constant):
+            if node.value in [None, True, False]:
+                return ast.NameConstant in types
+            elif isinstance(node.value, (int, float, complex)):
+                return ast.Num in types
+            elif isinstance(node.value, str):
+                return ast.Str in types
+            elif isinstance(node.value, bytes):
+                return ast.Bytes in types
+            elif node.value == Ellipsis:
+                return ast.Ellipsis in types
+            else:
+                raise RuntimeError('Unknown Constant value %r' % type(node.value))
+
+        return False
+
 class SuiteTransformer(NodeVisitor):
     """
     Transform suites of instructions
@@ -44,17 +94,39 @@ def __call__(self, node):
         return self.visit(node)
 
     def visit_ClassDef(self, node):
+        node.bases = [self.visit(b) for b in node.bases]
+
         node.body = self.suite(node.body, parent=node)
+        node.decorator_list = [self.visit(d) for d in node.decorator_list]
+
+        if hasattr(node, 'starargs') and node.starargs is not None:
+            node.starargs = self.visit(node.starargs)
+
+        if hasattr(node, 'kwargs') and node.kwargs is not None:
+            node.kwargs = self.visit(node.kwargs)
+
+        if hasattr(node, 'keywords'):
+            node.keywords = [self.visit(kw) for kw in node.keywords]
+
         return node
 
     def visit_FunctionDef(self, node):
+        node.args = self.visit(node.args)
         node.body = self.suite(node.body, parent=node)
+        node.decorator_list = [self.visit(d) for d in node.decorator_list]
+
+        if hasattr(node, 'returns') and node.returns is not None:
+            node.returns = self.visit(node.returns)
+
         return node
 
     def visit_AsyncFunctionDef(self, node):
         return self.visit_FunctionDef(node)
 
     def visit_For(self, node):
+        node.target = self.visit(node.target)
+        node.iter = self.visit(node.iter)
+
         node.body = self.suite(node.body, parent=node)
 
         if node.orelse:
@@ -63,15 +135,13 @@ def visit_For(self, node):
         return node
 
     def visit_AsyncFor(self, node):
-        node.body = self.suite(node.body, parent=node)
-
-        if node.orelse:
-            node.orelse = self.suite(node.orelse, parent=node)
-
-        return node
+        return self.visit_For(node)
 
     def visit_If(self, node):
+        node.test = self.visit(node.test)
+
         node.body = self.suite(node.body, parent=node)
+
         if node.orelse:
             node.orelse = self.suite(node.orelse, parent=node)
 
@@ -80,6 +150,8 @@ def visit_If(self, node):
     def visit_Try(self, node):
         node.body = self.suite(node.body, parent=node)
 
+        node.handlers = [self.visit(h) for h in node.handlers]
+
         if node.orelse:
             node.orelse = self.suite(node.orelse, parent=node)
 
@@ -89,6 +161,8 @@ def visit_Try(self, node):
         return node
 
     def visit_While(self, node):
+        node.test = self.visit(node.test)
+
         node.body = self.suite(node.body, parent=node)
 
         if node.orelse:
@@ -97,12 +171,20 @@ def visit_While(self, node):
         return node
 
     def visit_With(self, node):
+
+        if hasattr(node, 'items'):
+            node.items = [self.visit(i) for i in node.items]
+        else:
+            if node.context_expr:
+                node.context_expr = self.visit(node.context_expr)
+            if node.optional_vars:
+                node.optional_vars = self.visit(node.optional_vars)
+
         node.body = self.suite(node.body, parent=node)
         return node
 
     def visit_AsyncWith(self, node):
-        node.body = self.suite(node.body, parent=node)
-        return node
+        return self.visit_With(node)
 
     def visit_Module(self, node):
         node.body = self.suite(node.body, parent=node)
@@ -132,3 +214,29 @@ def generic_visit(self, node):
                 else:
                     setattr(node, field, new_node)
         return node
+
+    def add_child(self, child, parent, namespace=None):
+
+        def nearest_function_namespace(node):
+            """
+            Return the namespace node for the nearest function scope.
+
+            This could be itself.
+
+            :param node: The node to get the function namespace of
+            :type node: ast.Node
+            :rtype: ast.Node
+
+            """
+
+            if isinstance(node, (ast.FunctionDef, ast.Module)):
+                return node
+            if hasattr(ast, 'AsyncFunctionDef') and isinstance(node, ast.AsyncFunctionDef):
+                return node
+            return nearest_function_namespace(node.parent)
+
+        if namespace is None:
+            namespace = nearest_function_namespace(parent)
+
+        add_parent(child, parent=parent, namespace=namespace)
+        return child