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))
