Merge pull request #71 from dflook/comp-iter-namespace

Author: dflook

src/python_minifier/__init__.py

@@ -18,6 +18,7 @@
     allow_rename_locals,
     add_namespace,
 )
+
 from python_minifier.transforms.combine_imports import CombineImports
 from python_minifier.transforms.remove_annotations import RemoveAnnotations
 from python_minifier.transforms.remove_asserts import RemoveAsserts

src/python_minifier/expression_printer.py

@@ -16,7 +16,7 @@ def __init__(self):
         self.unicode_literals = False
 
         self.precedences = {
-            'Lambda': 2,  # Lamdda
+            'Lambda': 2,  # Lambda
             'IfExp': 3,  # IfExp
             'comprehension': 3.5,
             'Or': 4,  # BoolOp

src/python_minifier/rename/README.md

@@ -40,6 +40,7 @@ This is usually the closest parent namespace node. The exceptions are:
 - Function annotations are in the same namespace as their function.
 - Class decorator are in the same namespace as their class.
 - Class bases, keywords, starargs and kwargs are in the same namespace as their class.
+- The first iteration expression of a comprehension is in the same namespace as it's parent ListComp/SetComp/DictComp or GeneratorExp.
 
 ### Bind names
 

src/python_minifier/rename/mapper.py

@@ -89,6 +89,26 @@ def add_parent_to_classdef(classdef):
     for node in classdef.decorator_list:
         add_parent(node, parent=classdef, namespace=classdef.namespace)
 
+def add_parent_to_comprehension(node, namespace):
+    assert is_ast_node(node, (ast.GeneratorExp, 'SetComp', 'DictComp', 'ListComp'))
+
+    if hasattr(node, 'elt'):
+        add_parent(node.elt, parent=node, namespace=node)
+    elif hasattr(node, 'key'):
+        add_parent(node.key, parent=node, namespace=node)
+        add_parent(node.value, parent=node, namespace=node)
+
+    iter_namespace = namespace
+    for generator in node.generators:
+        generator.parent = node
+        generator.namespace = node
+
+        add_parent(generator.target, parent=generator, namespace=node)
+        add_parent(generator.iter, parent=generator, namespace=iter_namespace)
+        iter_namespace = node
+        for if_ in generator.ifs:
+            add_parent(if_, parent=generator, namespace=node)
+
 
 def add_parent(node, parent=None, namespace=None):
     """
@@ -114,6 +134,8 @@ def add_parent(node, parent=None, namespace=None):
 
         if is_ast_node(node, (ast.FunctionDef, 'AsyncFunctionDef')):
             add_parent_to_functiondef(node)
+        elif is_ast_node(node, (ast.GeneratorExp, 'SetComp', 'DictComp', 'ListComp')):
+            add_parent_to_comprehension(node, namespace=namespace)
         elif isinstance(node, ast.Lambda):
             add_parent_to_arguments(node.args, func=node)
             add_parent(node.body, parent=node, namespace=node)
@@ -125,13 +147,6 @@ def add_parent(node, parent=None, namespace=None):
 
         return
 
-    if isinstance(node, ast.comprehension):
-        add_parent(node.target, parent=node, namespace=namespace)
-        add_parent(node.iter, parent=node, namespace=namespace)
-        for if_ in node.ifs:
-            add_parent(if_, parent=node, namespace=namespace)
-        return
-
     if isinstance(node, ast.Global):
         namespace.global_names.update(node.names)
     if is_ast_node(node, 'Nonlocal'):

test/test_comprehension_rename.py

@@ -0,0 +1,96 @@
+import ast
+import sys
+
+import pytest
+
+from python_minifier import minify
+from python_minifier.ast_compare import compare_ast
+
+def test_listcomp_regression_2_7():
+    if sys.version_info >= (3, 0):
+        pytest.skip('ListComp doesn\'t create a new namespace in python < 3.0')
+
+    source = '''
+def f(pa):
+    return [pa.b for pa in pa.c]
+'''
+    expected = source
+    compare_ast(ast.parse(minify(source, rename_locals=True)), ast.parse(expected))
+
+def test_listcomp_regression():
+    if sys.version_info < (3, 0):
+        pytest.skip('ListComp creates a new namespace in python > 3.0')
+
+    source = '''
+def f(parentObject):
+    return [parentObject.b for parentObject in parentObject.c]
+'''
+    expected = '''
+def f(parentObject):
+    return[A.b for A in parentObject.c]
+'''
+    compare_ast(ast.parse(minify(source, rename_locals=True)), ast.parse(expected))
+
+
+def test_expression():
+    source = '''def test():
+    [x*y for x in range(10) for y in range(x, x+10)]
+'''
+
+    expected = '''def test():
+    [A*B for A in range(10) for B in range(A,A+10)]
+'''
+    compare_ast(ast.parse(minify(source, rename_locals=True)), ast.parse(expected))
+
+def test_generator_expression():
+    source = '''
+x=1
+def func():
+    return (x for x in x)
+'''
+
+    expected = '''
+x=1
+def func():
+    return (A for A in x)
+'''
+
+    compare_ast(ast.parse(minify(source, rename_locals=True)), ast.parse(expected))
+
+def test_generator_expression_multiple_for():
+    source = '''
+def func():
+    return (x for x in x for x in x)
+
+def func(long_name, another_long_name):
+    return (long_name for long_name, another_long_name in long_name for long_name in (long_name, another_long_name))
+'''
+
+    expected = '''
+def func():
+    return (A for A in x for A in A)
+
+def func(long_name, another_long_name):
+    return(A for A, B in long_name for A in (A,B))
+'''
+
+    compare_ast(ast.parse(minify(source, rename_locals=True)), ast.parse(expected))
+
+def test_generator_expression_nested_for():
+    source = '''
+def func():
+    return (a for a in (b for b in x) for c in c)
+
+def func(long_name):
+    return (a for a in (b for b in long_name) for c in c)
+'''
+
+    expected = '''
+def func():
+    return(A for A in (A for A in x) for B in B)
+
+def func(long_name):
+    return(A for A in (A for A in long_name) for B in B)
+'''
+
+    compare_ast(ast.parse(minify(source, rename_locals=True)), ast.parse(expected))