fixed renaming inside comprehensions, added binding in expressions

@@ -66,9 +66,33 @@ def unparse (st) :
 class Renamer (ast.NodeTransformer) :
 class Renamer (ast.NodeTransformer) :
     def __init__ (self, map_names) :
     def __init__ (self, map_names) :
         ast.NodeTransformer.__init__(self)
         ast.NodeTransformer.__init__(self)
-        self.map = map_names
+        self.map = [map_names]
+    def visit_ListComp (self, node) :
+        bind = self.map[-1].copy()
+        for comp in node.generators :
+            for name in getvars(comp.target) :
+                if name in bind :
+                    del bind[name]
+        self.map.append(bind)
+        node.elt = self.visit(node.elt)
+        self.map.pop(-1)
+        return node
+    def visit_SetComp (self, node) :
+        return self.visit_ListComp(node)
+    def visit_DictComp (self, node) :
+        bind = self.map[-1].copy()
+        for comp in node.generators :
+            for name in getvars(comp.target) :
+                if name in bind :
+                    del bind[name]
+        self.map.append(bind)
+        node.key = self.visit(node.key)
+        node.value = self.visit(node.value)
+        self.map.pop(-1)
+        return node
     def visit_Name (self, node) :
     def visit_Name (self, node) :
-        return ast.copy_location(ast.Name(id=self.map.get(node.id, node.id),
+        return ast.copy_location(ast.Name(id=self.map[-1].get(node.id,
+                                                              node.id),
                                           ctx=ast.Load()), node)
                                           ctx=ast.Load()), node)
 
 
 def rename (expr, map={}, **ren) :
 def rename (expr, map={}, **ren) :
@@ -77,12 +101,39 @@ def rename (expr, map={}, **ren) :
     '((t + y) < z)'
     '((t + y) < z)'
     >>> rename('x+y<z+f(3,t)', f='g', t='z', z='t')
     >>> rename('x+y<z+f(3,t)', f='g', t='z', z='t')
     '((x + y) < (t + g(3, z)))'
     '((x + y) < (t + g(3, z)))'
+    >>> rename('[x+y for x in range(3)]', x='z')
+    '[(x + y) for x in range(3)]'
+    >>> rename('[x+y for x in range(3)]', y='z')
+    '[(x + z) for x in range(3)]'
     """
     """
     map_names = dict(map)
     map_names = dict(map)
     map_names.update(ren)
     map_names.update(ren)
     transf = Renamer(map_names)
     transf = Renamer(map_names)
     return unparse(transf.visit(ast.parse(expr)))
     return unparse(transf.visit(ast.parse(expr)))
 
 
+class Binder (Renamer) :
+    def visit_Name (self, node) :
+        if node.id in self.map[-1] :
+            return self.map[-1][node.id]
+        else :
+            return node
+
+def bind (expr, map={}, **ren) :
+    """
+    >>> bind('x+y<z', x=ast.Num(n=2))
+    '((2 + y) < z)'
+    >>> bind('x+y<z', y=ast.Num(n=2))
+    '((x + 2) < z)'
+    >>> bind('[x+y for x in range(3)]', x=ast.Num(n=2))
+    '[(x + y) for x in range(3)]'
+    >>> bind('[x+y for x in range(3)]', y=ast.Num(n=2))
+    '[(x + 2) for x in range(3)]'
+    """
+    map_names = dict(map)
+    map_names.update(ren)
+    transf = Binder(map_names)
+    return unparse(transf.visit(ast.parse(expr)))
+
 if __name__ == "__main__" :
 if __name__ == "__main__" :
     import doctest
     import doctest
     doctest.testmod()
     doctest.testmod()

 from snakes.lang.ctlstar.parser import parse
 from snakes.lang.ctlstar.parser import parse
 from snakes.lang.ctlstar import asdl as ast
 from snakes.lang.ctlstar import asdl as ast
-from snakes.lang import getvars, unparse
+from snakes.lang import getvars, bind
 import _ast
 import _ast
 
 
 class SpecError (Exception) :
 class SpecError (Exception) :
@@ -19,40 +19,6 @@ def astcopy (node) :
             attr[name] = astcopy(value)
             attr[name] = astcopy(value)
     return node.__class__(**attr)
     return node.__class__(**attr)
 
 
-class Binder (ast.NodeTransformer) :
-    def __init__ (self, bind) :
-        ast.NodeTransformer.__init__(self)
-        self.bind = [bind]
-    def visit (self, node) :
-        return ast.NodeTransformer.visit(self, astcopy(node))
-    def visit_ListComp (self, node) :
-        """
-        >>> tree = ast.parse('x+y+[x+y+z for x, y in l]')
-        >>> unparse(Binder({'x':ast.Name('hello')}).visit(tree))
-        '((hello + y) + [((x + y) + z) for (x, y) in l])'
-        >>> unparse(Binder({'y':ast.Name('hello')}).visit(tree))
-        '((x + hello) + [((x + y) + z) for (x, y) in l])'
-        >>> unparse(Binder({'z':ast.Name('hello')}).visit(tree))
-        '((x + y) + [((x + y) + hello) for (x, y) in l])'
-        """
-        bind = self.bind[-1].copy()
-        for comp in node.generators :
-            for name in getvars(comp.target) :
-                if name in bind :
-                    del bind[name]
-        self.bind.append(bind)
-        node.elt = self.visit(node.elt)
-        self.bind.pop(-1)
-        return node
-    def visit_Name (self, node) :
-        if node.id in self.bind[-1] :
-            return astcopy(self.bind[-1][node.id])
-        else :
-            return astcopy(node)
-
-def bind (node, ctx) :
-    return Binder(ctx).visit(node)
-
 class Builder (object) :
 class Builder (object) :
     def __init__ (self, spec) :
     def __init__ (self, spec) :
         self.spec = spec
         self.spec = spec