def test_describe_function():
some_value = None
def f(a, b, c) -> Any:
"""hello world"""
return some_value
expected = FunctionDescription(
src=function_source(f),
ast=function_ast(f),
name='f',
qualname='test_describe_function.<locals>.f',
doc='hello world',
annotations={'return': Any},
defaults=f.__defaults__,
globals={'some_value': None},
referenced_modules=frozenset(),
module=f.__module__,
)
decomposed = describe(f)
assert (astunparse.dump(expected.ast) ==
astunparse.dump(decomposed.ast))
assert (expected._replace(ast=None) ==
decomposed._replace(ast=None))
def simplify_and_print(program: str):
tree = ast.parse(program)
if isinstance(tree, ast.Module) and len(tree.body) == 1:
tree = tree.body[0]
left = dump(tree)
tree = simplify(tree)
right = dump(tree)
if '\n' in left and '\n' in right:
print(juxtapose(left, right))
else:
print(left)
print(right)
def _transform(path: str, code: str,
target: CompilationTarget) -> Tuple[str, List[str]]:
"""Applies all transformation for passed target."""
debug(lambda: 'Compiling "{}"'.format(path))
dependencies = [] # type: List[str]
tree = ast.parse(code, path)
debug(lambda: 'Initial ast:\n{}'.format(dump(tree)))
for transformer in transformers:
if transformer.target < target:
debug(lambda: 'Skip transformer "{}"'.format(transformer.__name__))
continue
debug(lambda: 'Use transformer "{}"'.format(transformer.__name__))
working_tree = deepcopy(tree)
try:
result = transformer.transform(working_tree)
except SyntaxError as exc:
if isinstance(getattr(exc, 'ast_node', None), ast.AST):
if not getattr(exc.ast_node, 'lineno', None): # type: ignore
ast.fix_missing_locations(working_tree)
exc.lineno = getattr(exc.ast_node, 'lineno', 0) # type: ignore
exc.offset = getattr(exc.ast_node, 'col_offset',
-1) + 1 # type: ignore
else:
exc.lineno = exc.lineno or 0
exc.offset = exc.offset or 0
raise exc
except:
raise TransformationError(path, transformer, dump(tree),
format_exc())
if not result.tree_changed:
debug(lambda: 'Tree not changed')
continue
tree = working_tree
debug(lambda: 'Tree changed:\n{}'.format(dump(tree)))
dependencies.extend(result.dependencies)
try:
code = unparse(tree)
debug(lambda: 'Code changed:\n{}'.format(code))
except:
raise TransformationError(path, transformer, dump(tree),
format_exc())
# Disable E402 (moving imports to the top of the file) as it breaks.
code = fix_code(code,
options={'ignore': ['E226', 'E24', 'W50', 'W690', 'E402']})
return code, dependencies
def recurse_ast():
x = inspect.getsource(ast)
a = ast.parse(x)
code = astunparse.dump(a)
for x in xrange(0, 3):
print x
#print code
f = open("test_{x}.py".format(x=x), "w")
f.write(code)
f.close()
c = ast.parse(code)
code = astunparse.dump(c)
def new_method():
m2 = extract_ast(example_complex_class)
code = astunparse.dump(m2)
for x in xrange(0, 3):
print x
#print code
f = open("test_{x}.py".format(x=x), "w")
f.write(code)
f.close()
c = ast.parse(code)
code = astunparse.dump(c)
def print_ast(output_ast):
if arg_abstract_output:
try:
import astunparse
print astunparse.dump(output_ast)
except ImportError:
print ast.dump(output_ast)
else:
try:
import astunparse
print astunparse.unparse(output_ast)
except ImportError:
sys.exit("couldn't import 'astunparse'")
def new_method():
m2 = extract_ast(example_simple_function)
code = astunparse.dump(m2)
c2 = ast.parse(code)
code2 = astunparse.dump(c2)
rename = "from fast import Module, Expr, Call, Name, Load, keyword, List, Str, Num\n"
c3 = rename + "def eval_body():\n " + code2 + "\neval_body()"
f = open("test_eval_new_code_1.py", "w")
f.write(c3)
f.close()
c = compile(c3, '<string>', 'exec')
exec(c)
def visit_Attribute(self, node):
new_node = node
print('Attribute: ', astunparse.dump(new_node))
if self.isAttrRewritable(new_node.value):
new_node.attr = self.newName(new_node.attr)
self.visit(new_node.value)
return new_node
def old_method():
m = define_replace_method(
"extract_metavariable_methodname",
generate_args("extract_metavariable_list_element"),
Str(s='extract_metavariable_function_body'))
code = astunparse.dump(m)
print(code)
def check_ast_equals(a, b):
if isinstance(a, list):
assert isinstance(b, list)
a = ast.Module(body=a)
b = ast.Module(body=b)
if not ast.dump(a) == ast.dump(b):
differ = difflib.Differ()
a_ast = astunparse.dump(a).splitlines(keepends=True)
b_ast = astunparse.dump(b).splitlines(keepends=True)
a_src = astor.to_source(a).splitlines(keepends=True)
b_src = astor.to_source(b).splitlines(keepends=True)
diff = ''.join(differ.compare(a_src,
b_src)) if a_src != b_src else ''.join(
differ.compare(a_ast, b_ast))
return False, diff
return True, ''
def new_method():
m2 = extract_ast(example_simple_function)
code = astunparse.dump(m2)
c2 = ast.parse(code)
code2 = astunparse.dump(c2)
rename = "from fast import Module, Expr, Call, Name, Load, keyword, List, Str, Num\n"
c3 = rename + "def eval_body():\n "+code2 + "\neval_body()"
f = open ("test_eval_new_code_1.py","w")
f.write(c3)
f.close()
c = compile(c3, '<string>', 'exec')
exec(c)
def parse_code(filename, save_to_file):
original_code = open(filename, 'r').read()
root_node = parse(original_code)
if save_to_file:
filename2 = "parsed/parsed_" + ntpath.basename(filename)
with open(filename2, 'w') as file:
file.write(au.dump(root_node))
return root_node
def visit_With(self, node):
new_node = node
print('with: ', astunparse.dump(new_node))
for ind in range(len(new_node.items)):
self.visit(new_node.items[ind])
for ind in range(len(new_node.body)):
self.visit(new_node.body[ind])
return new_node
def test_dump_files_comparison(self):
"""Print the same data as other existing modules."""
for path in PATHS:
with open(path, 'r', encoding='utf-8') as py_file:
code = py_file.read()
untyped_tree = ast.parse(source=code, filename=path)
untyped_dump = astunparse.dump(untyped_tree)
tested_untyped_dump = typed_astunparse.dump(untyped_tree)
self.assertEqual(untyped_dump.splitlines(),
tested_untyped_dump.splitlines())
typed_tree = typed_ast.ast3.parse(source=code, filename=path)
bad_typed_dump = astunparse.dump(typed_tree)
for annotate_fields in [True, False]:
for include_attributes in [False, True]:
if include_attributes and not annotate_fields:
continue # behaviour differs from typed_ast
with self.assertRaises(TypeError):
_ = typed_ast.ast3.dump(
untyped_tree,
annotate_fields=annotate_fields,
include_attributes=include_attributes)
typed_dump = typed_ast.ast3.dump(
typed_tree,
annotate_fields=annotate_fields,
include_attributes=include_attributes)
tested_typed_dump = _postprocess_dump(
typed_astunparse.dump(
typed_tree,
annotate_fields=annotate_fields,
include_attributes=include_attributes))
if include_attributes:
# because of https://github.com/python/typed_ast/issues/23
self.assertEqual(typed_dump.replace(' ', ''),
tested_typed_dump.replace(' ', ''))
continue
self.assertNotEqual(untyped_dump, bad_typed_dump)
self.assertNotEqual(typed_dump, bad_typed_dump)
self.assertEqual(typed_dump, tested_typed_dump)
def visit_withitem(self, node):
new_node = node
print('withitem: ', astunparse.dump(new_node))
self.visit(new_node.context_expr)
if self.isrw_withitem(new_node.context_expr):
self.visit(new_node.optional_vars)
else:
self.names_built_in.add(new_node.optional_vars.id)
return new_node
def parse_code(filename, save_to_file):
original_code = open(filename, 'r').read()
root_node = parse(original_code)
if save_to_file:
filename = os.path.join('PARSED',
f'parsed_{ntpath.basename(filename)}')
file = open(filename, 'w')
file.write(au.dump(root_node))
file.close()
return root_node
def pprint_ast(mochi_file_name, ast_file_name=None, show_tokens=False):
"""Generate a nicly formatted AST from Mochi code.
"""
ast = translator.translate_file(mochi_file_name, show_tokens=show_tokens)
py_source = astunparse.dump(ast)
if not ast_file_name:
print(py_source)
else:
with open(ast_file_name, 'w') as fobj:
fobj.write(py_source)
def pprint_ast(mochi_file_name, ast_file_name=None, show_tokens=False):
"""Generate a nicly formatted AST from Mochi code.
"""
ast = translator.translate_file(mochi_file_name, show_tokens=show_tokens)
py_source = astunparse.dump(ast)
if not ast_file_name:
print(py_source)
else:
with open(ast_file_name, 'w') as fobj:
fobj.write(py_source)
def compile_and_run(filename, file_obj=None, debug=False, globals=None):
with file_obj or open(filename, 'r'):
ast_ = parse.parse(file_obj.read(), debug)
ast.fix_missing_locations(ast_)
with open('debug.py', 'w') as f:
f.write(astunparse.unparse(ast_))
with open('ast.out', 'w') as f:
f.write('\n------------ AST DUMP ------------\n')
f.write(astunparse.dump(ast_))
prog = compile(ast_, filename, 'exec')
exec(prog, globals or {})
def test_describe_function():
some_value = None
def f(a, b, c):
return some_value
expected = xun.functions.FunctionDescription(
src=xun.functions.function_source(f),
ast=xun.functions.function_ast(f),
name='f',
defaults=f.__defaults__,
globals={'some_value': None},
referenced_modules=frozenset(),
module=f.__module__,
)
decomposed = xun.functions.describe(f)
assert (astunparse.dump(expected.ast) ==
astunparse.dump(decomposed.ast))
assert (expected._replace(ast=None) ==
decomposed._replace(ast=None))
def parseprint(line: str):
tree = None
try:
tree = ast.parse(line)
except:
print("ast's parse raised an exception:", file=sys.stderr)
raise
try:
return print(astunparse.dump(ast.parse(line)))
except:
print("astunparse's dump raised an exception:", file=sys.stderr)
raise
def test_dump_files_comparison(self):
"""Print the same data as other existing modules."""
for path in PATHS:
with open(path, 'r', encoding='utf-8') as py_file:
code = py_file.read()
untyped_tree = ast.parse(source=code, filename=path)
untyped_dump = astunparse.dump(untyped_tree)
tested_untyped_dump = typed_astunparse.dump(untyped_tree)
self.assertEqual(untyped_dump.splitlines(), tested_untyped_dump.splitlines())
typed_tree = typed_ast.ast3.parse(source=code, filename=path)
bad_typed_dump = astunparse.dump(typed_tree)
for annotate_fields in [True, False]:
for include_attributes in [False, True]:
if include_attributes and not annotate_fields:
continue # behaviour differs from typed_ast
with self.assertRaises(TypeError):
_ = typed_ast.ast3.dump(
untyped_tree, annotate_fields=annotate_fields,
include_attributes=include_attributes)
typed_dump = typed_ast.ast3.dump(
typed_tree, annotate_fields=annotate_fields,
include_attributes=include_attributes)
tested_typed_dump = _postprocess_dump(typed_astunparse.dump(
typed_tree, annotate_fields=annotate_fields,
include_attributes=include_attributes))
if include_attributes:
# because of https://github.com/python/typed_ast/issues/23
self.assertEqual(
typed_dump.replace(' ', ''), tested_typed_dump.replace(' ', ''))
continue
self.assertNotEqual(untyped_dump, bad_typed_dump)
self.assertNotEqual(typed_dump, bad_typed_dump)
self.assertEqual(typed_dump, tested_typed_dump)
def _transform(path: str, code: str,
target: CompilationTarget) -> Tuple[str, List[str]]:
"""Applies all transformation for passed target."""
debug(lambda: 'Compiling "{}"'.format(path))
dependencies = [] # type: List[str]
tree = ast.parse(code, path)
debug(lambda: 'Initial ast:\n{}'.format(dump(tree)))
for transformer in transformers:
if transformer.target < target:
debug(lambda: 'Skip transformer "{}"'.format(transformer.__name__))
continue
debug(lambda: 'Use transformer "{}"'.format(transformer.__name__))
working_tree = deepcopy(tree)
try:
result = transformer.transform(working_tree)
except:
raise TransformationError(path, transformer, dump(tree),
format_exc())
if not result.tree_changed:
debug(lambda: 'Tree not changed')
continue
tree = working_tree
debug(lambda: 'Tree changed:\n{}'.format(dump(tree)))
dependencies.extend(result.dependencies)
try:
code = unparse(tree)
debug(lambda: 'Code changed:\n{}'.format(code))
except:
raise TransformationError(path, transformer, dump(tree),
format_exc())
return fix_code(code), dependencies
def find_variables_in_code(var_names, body, idx1):
elems = list()
idxs = list()
for idx, el in enumerate(body):
if idx > idx1:
names = find_names_in_code(au.dump(el))
for n in var_names:
if n in names:
elems.append(el)
names = []
idxs.append(idx)
for i in reversed(idxs):
del body[i]
return elems
def rewritefile(self, fileIn, fileOut):
file = open(fileIn, "r")
source = file.read()
unmodified = ast.parse(source)
source_r = astunparse.unparse(unmodified)
#print(source_r)
ast_unmodified = astunparse.dump(unmodified)
#print(ast_unmodified)
modified = self.visit(unmodified)
source_rob = astunparse.unparse(unmodified)
#print(source_rob)
with open(fileOut, "w") as text_file:
text_file.write("%s" % source_rob)
print(self.names_built_in)
joblib.dump(self.names_mapping, 'names_mapping.save')
def main():
fpath = argv[1]
with open(fpath, 'r') as f:
data = f.read()
auto_process_libs = [
'math',
'os.path',
]
auto_import_libs = [
*auto_process_libs,
'os',
]
options = Options(
debug=True,
imports=auto_import_libs,
eval_mod=auto_process_libs,
)
data = PreProcessor(options=options).preprocess(data)
data = Transpiler(options=options).transpile(data)
tree = parse(data)
tree = Generator(options=options).generate(tree)
tree = Optimizer(options=options).optimize(tree)
tree = Inliner(options=options).inline(tree)
tree = Importer( options=options).clean_imports(tree)
tree = UnusedRemover(options=options).remove_unused(tree)
code = Unparser.unparse(tree)
code = Minifier(options=options).minify(code)
log(dump(tree))
log(code.replace('\n', '\\n\n'))
path_parts = fpath.split('.')
out_path = join(getcwd(), 'dist',
''.join(path_parts[:-1])[1:] + '.py')
out_dir = dirname(out_path)
if not exists(out_dir):
mkdir(out_dir)
with open(out_path, 'w') as f:
f.write(code)
def visit_Call(self, node):
new_node = node
print('Call: ', astunparse.dump(node))
self.visit(new_node.func)
for ind in range(len(new_node.args)):
self.visit(new_node.args[ind])
for ind in range(len(new_node.keywords)):
if self.isKwFuncRewritable(new_node.func):
self.visit_keyword(new_node.keywords[ind])
else:
self.visit_keyword_noargRename(new_node.keywords[ind])
try:
self.visit(new_node.starargs)
except:
pass
try:
self.visit(new_node.kwargs)
except:
pass
return new_node
def top_level_visitor(node: concat.parse.TopLevelNode) -> ast.Module:
"""Converts a TopLevelNode to the top level of a Python module."""
statement = visitors.ref_visitor('statement')
word = visitors.ref_visitor('word')
body = list(All(Choice(statement, word)).visit(node))
statements = [
concat.astutils.statementfy(cast(Union[ast.stmt, ast.expr], child))
for child in body
]
module = ast.Module(body=statements)
ast.fix_missing_locations(module)
# debugging output
try:
with open('debug.py', 'w') as f:
f.write(astunparse.unparse(module))
with open('ast.out', 'w') as f:
f.write('------------ AST DUMP ------------\n')
f.write(astunparse.dump(module))
except UnicodeEncodeError:
pass
return module
def define_pattern2(methodname, thing, **kw):
args = []
newkw = {} # replace the string with variables
# create a set of loads for the variables
for k in kw:
v = kw[k]
p = create_param_name(v)
args.append(p)
newkw[k] = v
aast = extract_ast(thing)
ad = astunparse.dump(aast)
newast = ast.parse(ad)
newast2 = newast
body = do_replacements_objects(newast2, **newkw)
skip = body.body[0].value # why do we need this?
# we traverse the tree and remove the extra
# Module(body=[Expr(value=Call(
return define_replace_method(methodname, args, skip)
def define_pattern2(methodname, thing, **kw):
args = []
newkw = {} # replace the string with variables
# create a set of loads for the variables
for k in kw:
v = kw[k]
p = create_param_name(v)
args.append(p)
newkw[k] = v
aast = extract_ast(thing)
ad = astunparse.dump(aast)
newast = ast.parse(ad)
newast2 = newast
body = do_replacements_objects(newast2, **newkw)
skip = body.body[0].value # why do we need this?
# we traverse the tree and remove the extra
# Module(body=[Expr(value=Call(
return define_replace_method(methodname, args, skip)
def check_roundtrip(self, code1, filename="internal"):
ast_ = compile(str(code1), filename, "exec", ast.PyCF_ONLY_AST)
dump1 = astunparse.dump(ast_)
dump2 = ast.dump(ast_)
self.assertASTEqual(dump1, dump2)
def d(c):
print "dict",c
pprint.pprint(c._fields)
pprint.pprint(c.__dict__)
print(astunparse.unparse(c))
print(astunparse.dump(c))
def old_method():
m = define_replace_method("extract_metavariable_methodname",generate_args("extract_metavariable_list_element"),Str(s='extract_metavariable_function_body'))
code = astunparse.dump(m)
print(code)
def _str_to_node(string):
return _parse_expr(string)
def _import_module_code(name):
return ast.parse('import {0}; {0} = concatify({0})'.format(name)).body
def _set_line_info(p):
nodes = p[0] if isinstance(p[0], list) else [p[0]]
# print(p.lineno(1), p.lexpos(1))
for node in nodes:
node.lineno = p.lineno(1)
# print(node)
def parse(string, debug=0):
"""Parse a string in the Concat language."""
global debug_on
debug_on = debug
return ply.yacc.parse(string, lexer=lex.lexer, debug=debug)
ply.yacc.yacc()
if __name__ == '__main__':
while True:
tree = parse(input('Enter input >') + '\n')
print(astunparse.dump(tree))
def refine_args():
m2 = extract_ast(args)
code = astunparse.dump(m2)
name = extract_parent(replace(look_for('extract_metavariable_1'),"name"))
def new_method():
m2 = extract_ast(extract_metavariable_methodname)
code = astunparse.dump(m2)
print code
def dump(self, tree) -> str:
return astunparse.dump(tree)
def refine_body():
m2 = extract_ast(some_body)
code = astunparse.dump(m2)
print code
def check_roundtrip(self, code1, filename="internal", mode="exec"):
ast_ = compile(str(code1), filename, mode, ast.PyCF_ONLY_AST)
dump1 = astunparse.dump(ast_)
dump2 = ast.dump(ast_)
self.assertASTEqual(dump1, dump2)
NotEq, Pow, RShift, Slice, Store, Sub, Subscript, UAdd,
UnaryOp, USub, With, alias, keyword, withitem)
from typing import Any, List, Union
from astunparse import dump
from lark import Transformer
from lark.lexer import Token
from lark.tree import Tree
for elem in (Add, AnnAssign, Assign, Attribute, AugAssign, BitAnd, BitOr,
BitXor, Call, Compare, Constant, Dict, Div, Eq, FloorDiv, For,
FunctionDef, Gt, GtE, If, Load, LShift, Lt, LtE, MatMult, Mod,
Module, Mult, Name, NotEq, Pow, RShift, Store, Sub, keyword, With,
withitem, UnaryOp, Not, BinOp, alias, ImportFrom, UAdd, USub,
Index, Slice, Subscript):
setattr(elem, '__str__', lambda self: dump(self))
setattr(elem, '__repr__', lambda self: str(self))
class Treeify(Transformer):
DICT_OR_SET_MAKER = 'dictorsetmaker'
@staticmethod
def Constant(value: Any,
node: Union[Token, None] = None,
kind=None) -> Constant:
args = {
'value': value,
}
if node is not None:
import ast
import astunparse
print astunparse.dump(ast.parse('def func(self, dt): self.call(name)').body[0])
def ast_to_string(tree):
''' Return pretty-printed AST, as a string.
'''
return astunparse.dump(tree)
def t3st_unparse():
code = test_unparse_ast
print(astunparse.unparse(ast.parse(inspect.getsource(code))))
# get a pretty-printed dump of the AST
print(astunparse.dump(ast.parse(inspect.getsource(code))))
import sys
sys.path.append("/home/mdupont/experiments/astunparse/lib/")
import astunparse
import pprint
from ast import *
import ast
import inspect
def generate_metaclass(MetaClassVariable, NewBaseClassName, NewClassName):
return Module(body=[ClassDef(name=NewClassName, bases=[Name(id=NewBaseClassName, ctx=Load())], body=[Assign(targets=[Name(id='__metaclass__', ctx=Store())], value=Name(id=MetaClassVariable, ctx=Load()))], decorator_list=[])])
m = generate_metaclass("a","b","d")
#print(astunparse.unparse(m))
code = generate_metaclass
print(astunparse.unparse(ast.parse(inspect.getsource(code))))
# get a pretty-printed dump of the AST
print(astunparse.dump(ast.parse(inspect.getsource(code))))
def test_from_import_star(self):
"""Test for 'from module import *'."""
expected = textwrap.dedent("""
Module(body=[
ImportFrom(
module='concat.libconcat',
names=[alias(
name='*',
asname=None)],
level=0),
Assign(
targets=[Name(
id='itertools',
ctx=Store())],
value=Call(
func=Name(
id='import_and_convert',
ctx=Load()),
args=[Str(s='itertools')],
keywords=[])),
Assign(
targets=[Name(
id='accumulate',
ctx=Store())],
value=Attribute(
value=Name(
id='itertools',
ctx=Load()),
attr='accumulate',
ctx=Load())),
Assign(
targets=[Name(
id='chain',
ctx=Store())],
value=Attribute(
value=Name(
id='itertools',
ctx=Load()),
attr='chain',
ctx=Load())),
Assign(
targets=[Name(
id='combinations',
ctx=Store())],
value=Attribute(
value=Name(
id='itertools',
ctx=Load()),
attr='combinations',
ctx=Load())),
Assign(
targets=[Name(
id='combinations_with_replacement',
ctx=Store())],
value=Attribute(
value=Name(
id='itertools',
ctx=Load()),
attr='combinations_with_replacement',
ctx=Load())),
Assign(
targets=[Name(
id='compress',
ctx=Store())],
value=Attribute(
value=Name(
id='itertools',
ctx=Load()),
attr='compress',
ctx=Load())),
Assign(
targets=[Name(
id='count',
ctx=Store())],
value=Attribute(
value=Name(
id='itertools',
ctx=Load()),
attr='count',
ctx=Load())),
Assign(
targets=[Name(
id='cycle',
ctx=Store())],
value=Attribute(
value=Name(
id='itertools',
ctx=Load()),
attr='cycle',
ctx=Load())),
Assign(
targets=[Name(
id='dropwhile',
ctx=Store())],
value=Attribute(
value=Name(
id='itertools',
ctx=Load()),
attr='dropwhile',
ctx=Load())),
Assign(
targets=[Name(
id='filterfalse',
ctx=Store())],
value=Attribute(
value=Name(
id='itertools',
ctx=Load()),
attr='filterfalse',
ctx=Load())),
Assign(
targets=[Name(
id='groupby',
ctx=Store())],
value=Attribute(
value=Name(
id='itertools',
ctx=Load()),
attr='groupby',
ctx=Load())),
Assign(
targets=[Name(
id='islice',
ctx=Store())],
value=Attribute(
value=Name(
id='itertools',
ctx=Load()),
attr='islice',
ctx=Load())),
Assign(
targets=[Name(
id='permutations',
ctx=Store())],
value=Attribute(
value=Name(
id='itertools',
ctx=Load()),
attr='permutations',
ctx=Load())),
Assign(
targets=[Name(
id='product',
ctx=Store())],
value=Attribute(
value=Name(
id='itertools',
ctx=Load()),
attr='product',
ctx=Load())),
Assign(
targets=[Name(
id='repeat',
ctx=Store())],
value=Attribute(
value=Name(
id='itertools',
ctx=Load()),
attr='repeat',
ctx=Load())),
Assign(
targets=[Name(
id='starmap',
ctx=Store())],
value=Attribute(
value=Name(
id='itertools',
ctx=Load()),
attr='starmap',
ctx=Load())),
Assign(
targets=[Name(
id='takewhile',
ctx=Store())],
value=Attribute(
value=Name(
id='itertools',
ctx=Load()),
attr='takewhile',
ctx=Load())),
Assign(
targets=[Name(
id='tee',
ctx=Store())],
value=Attribute(
value=Name(
id='itertools',
ctx=Load()),
attr='tee',
ctx=Load())),
Assign(
targets=[Name(
id='zip_longest',
ctx=Store())],
value=Attribute(
value=Name(
id='itertools',
ctx=Load()),
attr='zip_longest',
ctx=Load()))])
""").strip()
actual = astunparse.dump(parse('from itertools import *\n')).strip()
self.assertEqual(actual, expected)