1. Victor Stinner
  2. astoptimizer

Commits

Victor Stinner  committed d6ac309

Fix Optimizer.optimize_comprehension() for dict(generator)

  • Participants
  • Parent commits c2c4718
  • Branches default

Comments (0)

Files changed (1)

File astoptimizer/optimizer.py

View file
  • Ignore whitespace
                 return
             key = target.elts[0].id
             value = target.elts[1].id
-            if not (isinstance(node.key, ast.Name)
+            if isinstance(node, ast.DictComp):
+                simple_gen = (
+                    isinstance(node.key, ast.Name)
                     and isinstance(node.value, ast.Name)
                     and node.key.id == key
-                    and node.value.id == value):
-                # {value: key for key, value in iterable}
+                    and node.value.id == value)
+            else:
+                simple_gen = (
+                    isinstance(node.elt, ast.Tuple)
+                    and len(node.elt.elts) == 2
+                    and isinstance(node.elt.elts[0], ast.Name)
+                    and node.elt.elts[0].id == key
+                    and isinstance(node.elt.elts[1], ast.Name)
+                    and node.elt.elts[1].id == value)
+            if not simple_gen:
+                # Example: {value: key for key, value in iterable}
                 if is_generator:
                     new_iter = self.optimize_generator(generator.iter)
                 else: