def tree_embeds(t, u, leaves=()): ''' ... >>> tree_embeds([1], [2]) True >>> tree_embeds([1], [1,2]) False >>> tree_embeds([1], [[1,2]]) True >>> tree_embeds([1,[2,3]], [[True, False], 'ab']) True >>> tree_embeds([1,[2,3]], [[True, False], 'ab'], leaves=(str,)) False >>> tree_embeds([], [1,2,3]) False >>> tree_embeds(1, [1,2,3]) True ''' # Pass leaf types around, avoiding infinite regress on strings _is_fork = partial(is_fork, leaves=leaves) if isinstance(t, basestring) or isinstance(u, basestring): leaves += (basestring,) _tree_embeds = partial(tree_embeds, leaves=leaves) if not _is_fork(t): return True else: return (_is_fork(u) and len(list(t)) == len(list(u)) and all(_tree_embeds(n,m) for n,m in zip(t,u)))
def flatten(tree, leaves=()): ''' Flatten a tree, that is, recursively concatenate nested sequences. >>> flatten([1, [[2, 3], 4, 5], [[[[6]]]]]) [1, 2, 3, 4, 5, 6] >>> mixed = [1, [[2], 3], (4, 5), [(6,)], 'abc'] >>> flatten(mixed) [1, 2, 3, 4, 5, 6, 'a', 'b', 'c'] >>> flatten(mixed, leaves=(str,)) [1, 2, 3, 4, 5, 6, 'abc'] >>> flatten(mixed, leaves=(str, tuple)) [1, 2, 3, (4, 5), (6,), 'abc'] >>> flatten([]) [] >>> flatten(1) [1] ''' # Pass leaf types around, avoiding infinite regress on strings _is_fork = partial(is_fork, leaves=leaves) if isinstance(tree, basestring): leaves += (basestring, ) _flatten = partial(flatten, leaves=leaves) if not _is_fork(tree): return [tree] else: return concat(_flatten(n) for n in tree)
def flatten(tree, leaves=()): ''' Flatten a tree, that is, recursively concatenate nested sequences. >>> flatten([1, [[2, 3], 4, 5], [[[[6]]]]]) [1, 2, 3, 4, 5, 6] >>> mixed = [1, [[2], 3], (4, 5), [(6,)], 'abc'] >>> flatten(mixed) [1, 2, 3, 4, 5, 6, 'a', 'b', 'c'] >>> flatten(mixed, leaves=(str,)) [1, 2, 3, 4, 5, 6, 'abc'] >>> flatten(mixed, leaves=(str, tuple)) [1, 2, 3, (4, 5), (6,), 'abc'] >>> flatten([]) [] >>> flatten(1) [1] ''' # Pass leaf types around, avoiding infinite regress on strings _is_fork = partial(is_fork, leaves=leaves) if isinstance(tree, basestring): leaves += (basestring,) _flatten = partial(flatten, leaves=leaves) if not _is_fork(tree): return [tree] else: return concat(_flatten(n) for n in tree)
def tree_embeds(t, u, leaves=()): ''' ... >>> tree_embeds([1], [2]) True >>> tree_embeds([1], [1,2]) False >>> tree_embeds([1], [[1,2]]) True >>> tree_embeds([1,[2,3]], [[True, False], 'ab']) True >>> tree_embeds([1,[2,3]], [[True, False], 'ab'], leaves=(str,)) False >>> tree_embeds([], [1,2,3]) False >>> tree_embeds(1, [1,2,3]) True ''' # Pass leaf types around, avoiding infinite regress on strings _is_fork = partial(is_fork, leaves=leaves) if isinstance(t, basestring) or isinstance(u, basestring): leaves += (basestring, ) _tree_embeds = partial(tree_embeds, leaves=leaves) if not _is_fork(t): return True else: return (_is_fork(u) and len(list(t)) == len(list(u)) and all(_tree_embeds(n, m) for n, m in zip(t, u)))
def _transform_children(self, node): if not isinstance(node, Node): raise ValueError('Expected a Node, got %r' % node) # In these trees, strings and Nodes are leaves leaves = (str, Node) tree_map = partial(tree.tree_map, leaves=leaves) children = tree_map(partial(self.transform), node.getChildrenTree()) return node.__class__(*children)
def tree_map(f, tree, leaves=()): ''' Map a function over the leaves of a tree. >>> tree_map(str, [1,[2,3]]) ['1', ['2', '3']] >>> tree_map(len, ['foo', ['x', 'asdf']]) [[1, 1, 1], [[1], [1, 1, 1, 1]]] >>> tree_map(len, ['foo', ['x', 'asdf']], leaves=(str,)) [3, [1, 4]] ''' # Pass leaf types around, avoiding infinite regress on strings _is_fork = partial(is_fork, leaves=leaves) if isinstance(tree, basestring): leaves += (basestring,) _tree_map = partial(tree_map, leaves=leaves) if not _is_fork(tree): return f(tree) else: seq = isinstance(tree, tuple) and tuple or list return seq([ _tree_map(f,n) for n in tree ])
def tree_map(f, tree, leaves=()): ''' Map a function over the leaves of a tree. >>> tree_map(str, [1,[2,3]]) ['1', ['2', '3']] >>> tree_map(len, ['foo', ['x', 'asdf']]) [[1, 1, 1], [[1], [1, 1, 1, 1]]] >>> tree_map(len, ['foo', ['x', 'asdf']], leaves=(str,)) [3, [1, 4]] ''' # Pass leaf types around, avoiding infinite regress on strings _is_fork = partial(is_fork, leaves=leaves) if isinstance(tree, basestring): leaves += (basestring, ) _tree_map = partial(tree_map, leaves=leaves) if not _is_fork(tree): return f(tree) else: seq = isinstance(tree, tuple) and tuple or list return seq([_tree_map(f, n) for n in tree])
def tree_zip(*trees, **kw): ''' Zip recursively. Returns nested lists with tuples at the bottom. >>> tree_zip([1,[2,3]], [5,[6,7]]) [(1, 5), [(2, 6), (3, 7)]] >>> tree_zip([1,[2,3]], [5,[6,7]], ['a',['b','c']]) [(1, 5, 'a'), [(2, 6, 'b'), (3, 7, 'c')]] >>> tree_zip('foo', 'bar') [('f', 'b'), ('o', 'a'), ('o', 'r')] >>> tree_zip(['foo'], ['bar']) [[('f', 'b'), ('o', 'a'), ('o', 'r')]] >>> tree_zip(['foo'], ['bar'], leaves=(str,)) [('foo', 'bar')] >>> tree_zip('foo', ['bar']) [('f', 'bar')] >>> tree_zip('foo', ['bar'], leaves=(str,)) ('foo', ['bar']) >>> tree_zip(1, 2) (1, 2) >>> tree_zip([1,[2,3]], [5,6,7,8], [['a','b'],['c','d']]) [(1, 5, ['a', 'b']), ([2, 3], 6, ['c', 'd'])] ''' # (Default arguments) leaves = kw.get('leaves', ()) # Pass leaf types around, avoiding infinite regress on strings _is_fork = partial(is_fork, leaves=leaves) if any(isinstance(t, basestring) for t in trees): leaves += (basestring,) _tree_zip = partial(tree_zip, leaves=leaves) if not all(_is_fork(t) for t in trees): return trees else: return [ _tree_zip(*neighbors) for neighbors in zip(*trees) ]
def tree_zip(*trees, **kw): ''' Zip recursively. Returns nested lists with tuples at the bottom. >>> tree_zip([1,[2,3]], [5,[6,7]]) [(1, 5), [(2, 6), (3, 7)]] >>> tree_zip([1,[2,3]], [5,[6,7]], ['a',['b','c']]) [(1, 5, 'a'), [(2, 6, 'b'), (3, 7, 'c')]] >>> tree_zip('foo', 'bar') [('f', 'b'), ('o', 'a'), ('o', 'r')] >>> tree_zip(['foo'], ['bar']) [[('f', 'b'), ('o', 'a'), ('o', 'r')]] >>> tree_zip(['foo'], ['bar'], leaves=(str,)) [('foo', 'bar')] >>> tree_zip('foo', ['bar']) [('f', 'bar')] >>> tree_zip('foo', ['bar'], leaves=(str,)) ('foo', ['bar']) >>> tree_zip(1, 2) (1, 2) >>> tree_zip([1,[2,3]], [5,6,7,8], [['a','b'],['c','d']]) [(1, 5, ['a', 'b']), ([2, 3], 6, ['c', 'd'])] ''' # (Default arguments) leaves = kw.get('leaves', ()) # Pass leaf types around, avoiding infinite regress on strings _is_fork = partial(is_fork, leaves=leaves) if any(isinstance(t, basestring) for t in trees): leaves += (basestring, ) _tree_zip = partial(tree_zip, leaves=leaves) if not all(_is_fork(t) for t in trees): return trees else: return [_tree_zip(*neighbors) for neighbors in zip(*trees)]
def tree_shape(tree, leaves=()): ''' The shape of a tree expressed as nested tuples of nothing. >>> tree_shape([1,[2,3]]) == tree_shape([True, 'ab']) True >>> tree_shape([1,[2,3]]) ((), ((), ())) >>> tree_shape(['ab', ['c', 'd']]) (((), ()), (((),), ((),))) >>> tree_shape(['ab', ['c', 'd']], leaves=(str,)) ((), ((), ())) >>> tree_shape(1) () ''' # Pass leaf types around, avoiding infinite regress on strings _is_fork = partial(is_fork, leaves=leaves) if isinstance(tree, basestring): leaves += (basestring,) _tree_shape = partial(tree_shape, leaves=leaves) if not _is_fork(tree): return () else: return tuple(_tree_shape(n) for n in tree)
def tree_shape(tree, leaves=()): ''' The shape of a tree expressed as nested tuples of nothing. >>> tree_shape([1,[2,3]]) == tree_shape([True, 'ab']) True >>> tree_shape([1,[2,3]]) ((), ((), ())) >>> tree_shape(['ab', ['c', 'd']]) (((), ()), (((),), ((),))) >>> tree_shape(['ab', ['c', 'd']], leaves=(str,)) ((), ((), ())) >>> tree_shape(1) () ''' # Pass leaf types around, avoiding infinite regress on strings _is_fork = partial(is_fork, leaves=leaves) if isinstance(tree, basestring): leaves += (basestring, ) _tree_shape = partial(tree_shape, leaves=leaves) if not _is_fork(tree): return () else: return tuple(_tree_shape(n) for n in tree)
def transformAssign(self, node): # Only transform constant assignments if not is_const(node.expr): return self._transform_children(node) # '__foo' names come from each lhs name at the head of the assignment: # 'l = a,b = 0,1' -> 'l = a,b = __a,__b', { '__a':0, '__b':1 } # # Unpack the rhs enough to map each '__foo' name to a value. # 'a,l = 0,[1,2]' -> 'a,l = __a,__l', { '__a':0, '__l':[1,2] } # 'a,l = 2' -> SyntaxError def lvalue_name(node): 'Construct a "magic" name to represent an l-value.' prefix = sep = '__' dot = '_' if isinstance(node, AssName): return prefix + node.name elif isinstance(node, AssAttr): name = node.attrname expr = node.expr while isinstance(expr, Getattr): name = sep.join([expr.attrname, name]) expr = expr.expr if isinstance(expr, Name): expr_name = expr.name else: expr_name = dot return prefix + sep.join([expr_name, name]) # In these trees, strings and tuples are leaves leaves = (str, tuple, AssAttr, AssName) tree_zip = partial(tree.tree_zip, leaves=leaves) flatten = partial(tree.flatten, leaves=leaves) tree_embeds = partial(tree.tree_embeds, leaves=leaves) # Grab the (right-most) lhs and the rhs lhs, rhs = node.nodes[-1], node.expr # Associate constants with l-value names if not tree_embeds(lhs, rhs): raise SyntaxError('Not enough r-values to unpack: %s' % node) zipped = flatten(tree_zip(lhs, rhs)) const_ast_for = map_keys(lambda v: lvalue_name(v), dict(zipped)) # Gather name<->const mappings for names we haven't seen before name_for = {} for name in const_ast_for.keys(): if name not in self.const_for.keys(): self.const_for[name] = eval_ast(Expression(const_ast_for[name])) assert const_ast_for[name] not in name_for name_for[const_ast_for[name]] = name class C(Transformer): def transform(self, node): if isinstance(node, Node) and node in name_for.keys(): return Name(name_for[node]) else: return super(C, self).transform(node) return Assign(node.nodes, C().transform(rhs))