def test_is_shebang_comment(self):
"""
Tests whether the libfuturize.fixer_util.is_shebang_comment() function is working
"""
node = FromImport(u'math', [Leaf(token.NAME, u'cos', prefix=" ")])
node.prefix = u'#!/usr/bin/env python\n'
self.assertTrue(is_shebang_comment(node))
def test_is_shebang_comment(self):
"""
Tests whether the libfuturize.fixer_util.is_shebang_comment() function is working
"""
node = FromImport(u'math', [Leaf(token.NAME, u'cos', prefix=" ")])
node.prefix = u'#!/usr/bin/env python\n'
self.assertTrue(is_shebang_comment(node))
def transform(self, node, results):
imp = results.get('imp')
if node.type == syms.import_from:
# Some imps are top-level (e.g., 'import ham')
# some are first level (e.g., 'import ham.eggs')
# some are third level (e.g., 'import ham.eggs as spam')
# Hence, the loop
while not hasattr(imp, 'value'):
imp = imp.children[0]
if imp.value.startswith("ui_"):
imp.value = u"." + imp.value
imp.changed()
else:
first = imp
if isinstance(imp, Node):
first = next(imp.leaves())
if not isinstance(first, Leaf):
return
if not first.value.startswith("ui_"):
return
new = FromImport(".", [imp])
new.prefix = node.prefix
return new
def transform(self, node, results):
if self.skip:
return
imp = results['imp']
if node.type == syms.import_from:
# Some imps are top-level (eg: 'import ham')
# some are first level (eg: 'import ham.eggs')
# some are third level (eg: 'import ham.eggs as spam')
# Hence, the loop
while not hasattr(imp, 'value'):
imp = imp.children[0]
if self.probably_a_local_import(imp.value):
imp.value = "." + imp.value
imp.changed()
else:
have_local = False
have_absolute = False
for mod_name in traverse_imports(imp):
if self.probably_a_local_import(mod_name):
have_local = True
else:
have_absolute = True
if have_absolute:
if have_local:
# We won't handle both sibling and absolute imports in the
# same statement at the moment.
self.warning(node, "absolute and local imports together")
return
new = FromImport(".", [imp])
new.prefix = node.prefix
return new
def new_future_import(self, old):
new = FromImport("__future__",
[Name("absolute_import", prefix=" "), Comma(),
Name("division", prefix=" "), Comma(),
Name("print_function", prefix=" ")])
if old is not None:
new.prefix = old.prefix
return new
def new_future_import(self, old):
new = FromImport("__future__", [
Name("absolute_import", prefix=" "),
Comma(),
Name("division", prefix=" "),
Comma(),
Name("print_function", prefix=" ")
])
if old is not None:
new.prefix = old.prefix
return new
def test_is_shebang_comment(self):
"""
Tests whether the fixer_util.is_encoding_comment() function is working.
"""
shebang_comments = [u"#!/usr/bin/env python\n" u"#!/usr/bin/python2\n", u"#! /usr/bin/python3\n"]
not_shebang_comments = [u"# I saw a giant python\n", u"# I have never seen a python2\n"]
for comment in shebang_comments:
node = FromImport(u"math", [Leaf(token.NAME, u"cos", prefix=" ")])
node.prefix = comment
self.assertTrue(is_shebang_comment(node))
for comment in not_shebang_comments:
node = FromImport(u"math", [Leaf(token.NAME, u"cos", prefix=" ")])
node.prefix = comment
self.assertFalse(is_shebang_comment(node))
def future_import2(feature, node):
"""
An alternative to future_import() which might not work ...
"""
root = find_root(node)
if does_tree_import(u"__future__", feature, node):
return
insert_pos = 0
for idx, node in enumerate(root.children):
if node.type == syms.simple_stmt and node.children and \
node.children[0].type == token.STRING:
insert_pos = idx + 1
break
for thing_after in root.children[insert_pos:]:
if thing_after.type == token.NEWLINE:
insert_pos += 1
continue
prefix = thing_after.prefix
thing_after.prefix = u""
break
else:
prefix = u""
import_ = FromImport(u"__future__",
[Leaf(token.NAME, feature, prefix=u" ")])
children = [import_, Newline()]
root.insert_child(insert_pos,
Node(syms.simple_stmt, children, prefix=prefix))
def future_import(feature, node):
"""
This seems to work
"""
root = find_root(node)
if does_tree_import(u"__future__", feature, node):
return
for idx, node in enumerate(root.children):
if node.type == syms.simple_stmt and \
len(node.children) > 0 and node.children[0].type == token.STRING:
# skip over docstring
continue
names = check_future_import(node)
if not names:
# not a future statement; need to insert before this
break
if feature in names:
# already imported
return
import_ = FromImport(u'__future__',
[Leaf(token.NAME, feature, prefix=" ")])
children = [import_, Newline()]
root.insert_child(idx, Node(syms.simple_stmt, children))
def touch_import_top(package, name_to_import, node):
"""Works like `does_tree_import` but adds an import statement at the
top if it was not imported (but below any __future__ imports).
Calling this multiple times adds them in reverse order.
Based on lib2to3.fixer_util.touch_import()
"""
root = find_root(node)
if does_tree_import(package, name_to_import, root):
return
# Look for __future__ imports and insert below them
found = False
for name in [
'absolute_import', 'division', 'print_function', 'unicode_literals'
]:
if does_tree_import('__future__', name, root):
found = True
break
if found:
# At least one __future__ import. We want to loop until we've seen them
# all.
start, end = None, None
for idx, node in enumerate(root.children):
if check_future_import(node):
start = idx
# Start looping
idx2 = start
while node:
node = node.next_sibling
idx2 += 1
if not check_future_import(node):
end = idx2
break
break
assert start is not None
assert end is not None
insert_pos = end
else:
# No __future__ imports
for idx, node in enumerate(root.children):
if node.type == syms.simple_stmt: # and node.children and node.children[0].type == token.STRING):
break
insert_pos = idx
if package is None:
import_ = Node(syms.import_name, [
Leaf(token.NAME, u"import"),
Leaf(token.NAME, name_to_import, prefix=u" ")
])
else:
import_ = FromImport(package,
[Leaf(token.NAME, name_to_import, prefix=u" ")])
children = [import_, Newline()]
root.insert_child(insert_pos, Node(syms.simple_stmt, children))
def test_is_encoding_comment(self):
"""
Tests whether the fixer_util.is_encoding_comment() function is working.
"""
encoding_comments = [
u"# coding: utf-8",
u"# encoding: utf-8",
u"# -*- coding: latin-1 -*-",
u"# vim: set fileencoding=iso-8859-15 :",
]
not_encoding_comments = [u"# We use the file encoding utf-8", u"coding = 'utf-8'", u"encoding = 'utf-8'"]
for comment in encoding_comments:
node = FromImport(u"math", [Leaf(token.NAME, u"cos", prefix=" ")])
node.prefix = comment
self.assertTrue(is_encoding_comment(node))
for comment in not_encoding_comments:
node = FromImport(u"math", [Leaf(token.NAME, u"cos", prefix=" ")])
node.prefix = comment
self.assertFalse(is_encoding_comment(node))
def add_globals(self, node):
"""Add required globals to the root of node. Idempotent."""
if self.added_pyi_globals:
return
# TODO: get rid of this -- added to prevent adding .parsed_pyi.top_lines every time
# we annotate a different function in the same file, but can break when we run the tool
# twice on the same file. Have to do something like what touch_import does.
self.added_pyi_globals = True
imports, top_lines = self.parsed_pyi.imports, self.parsed_pyi.top_lines
# Copy imports if not already present
for pkg, names in imports:
if names is None:
# TODO: do ourselves, touch_import puts stuff above license headers
touch_import(None, pkg, node) # == 'import pkg'
else:
for name in names:
touch_import(pkg, name, node)
root = find_root(node)
import_idx = [
idx for idx, node in enumerate(root.children)
if self.import_pattern.match(node)
]
if import_idx:
future_insert_pos = import_idx[0]
top_insert_pos = import_idx[-1] + 1
else:
future_insert_pos = top_insert_pos = 0
# first string (normally docstring)
for idx, node in enumerate(root.children):
if (node.type == syms.simple_stmt and node.children
and node.children[0].type == token.STRING):
future_insert_pos = top_insert_pos = idx + 1
break
top_lines = '\n'.join(top_lines)
top_lines = Util.parse_string(top_lines) # strips some newlines
for offset, node in enumerate(top_lines.children[:-1]):
root.insert_child(top_insert_pos + offset, node)
# touch_import doesn't do proper order for __future__
pkg = '__future__'
future_imports = [
n for n in self.future_imports
if not does_tree_import(pkg, n, root)
]
for offset, name in enumerate(future_imports):
node = FromImport(pkg, [Leaf(token.NAME, name, prefix=" ")])
node = Node(syms.simple_stmt, [node, Newline()])
root.insert_child(future_insert_pos + offset, node)
def transform(self, node, results):
"""
Copied from FixImport.transform(), but with this line added in
any modules that had implicit relative imports changed:
from __future__ import absolute_import"
"""
if self.skip:
return
imp = results['imp']
if node.type == syms.import_from:
# Some imps are top-level (eg: 'import ham')
# some are first level (eg: 'import ham.eggs')
# some are third level (eg: 'import ham.eggs as spam')
# Hence, the loop
while not hasattr(imp, 'value'):
imp = imp.children[0]
if self.probably_a_local_import(imp.value):
imp.value = u"." + imp.value
imp.changed()
future_import(u"absolute_import", node)
else:
have_local = False
have_absolute = False
for mod_name in traverse_imports(imp):
if self.probably_a_local_import(mod_name):
have_local = True
else:
have_absolute = True
if have_absolute:
if have_local:
# We won't handle both sibling and absolute imports in the
# same statement at the moment.
self.warning(node, "absolute and local imports together")
return
new = FromImport(u".", [imp])
new.prefix = node.prefix
future_import(u"absolute_import", node)
return new
def future_import(feature, node):
"""
This seems to work
"""
root = find_root(node)
if does_tree_import(u"__future__", feature, node):
return
# Look for a shebang or encoding line
shebang_encoding_idx = None
for idx, node in enumerate(root.children):
# Is it a shebang or encoding line?
if is_shebang_comment(node) or is_encoding_comment(node):
shebang_encoding_idx = idx
if node.type == syms.simple_stmt and \
len(node.children) > 0 and node.children[0].type == token.STRING:
# skip over docstring
continue
names = check_future_import(node)
if not names:
# not a future statement; need to insert before this
break
if feature in names:
# already imported
return
import_ = FromImport(u'__future__',
[Leaf(token.NAME, feature, prefix=" ")])
if shebang_encoding_idx == 0 and idx == 0:
# If this __future__ import would go on the first line,
# detach the shebang / encoding prefix from the current first line.
# and attach it to our new __future__ import node.
import_.prefix = root.children[0].prefix
root.children[0].prefix = u''
# End the __future__ import line with a newline and add a blank line
# afterwards:
children = [import_, Newline()]
root.insert_child(idx, Node(syms.simple_stmt, children))
def future_import(feature, node):
"""
This seems to work
"""
root = find_root(node)
if does_tree_import(u"__future__", feature, node):
return
# Look for a shebang or encoding line
shebang_encoding_idx = None
for idx, node in enumerate(root.children):
# Is it a shebang or encoding line?
if is_shebang_comment(node) or is_encoding_comment(node):
shebang_encoding_idx = idx
if node.type == syms.simple_stmt and \
len(node.children) > 0 and node.children[0].type == token.STRING:
# skip over docstring
continue
names = check_future_import(node)
if not names:
# not a future statement; need to insert before this
break
if feature in names:
# already imported
return
import_ = FromImport(u'__future__', [Leaf(token.NAME, feature, prefix=" ")])
if shebang_encoding_idx == 0 and idx == 0:
# If this __future__ import would go on the first line,
# detach the shebang / encoding prefix from the current first line.
# and attach it to our new __future__ import node.
import_.prefix = root.children[0].prefix
root.children[0].prefix = u''
# End the __future__ import line with a newline and add a blank line
# afterwards:
children = [import_ , Newline()]
root.insert_child(idx, Node(syms.simple_stmt, children))
def add_future(node, symbol):
root = find_root(node)
for idx, node in enumerate(root.children):
if node.type == syms.simple_stmt and \
len(node.children) > 0 and node.children[0].type == token.STRING:
# skip over docstring
continue
names = check_future_import(node)
if not names:
# not a future statement; need to insert before this
break
if symbol in names:
# already imported
return
import_ = FromImport('__future__', [Leaf(token.NAME, symbol, prefix=" ")])
children = [import_, Newline()]
root.insert_child(idx, Node(syms.simple_stmt, children))
def test_is_shebang_comment(self):
"""
Tests whether the fixer_util.is_encoding_comment() function is working.
"""
shebang_comments = [u'#!/usr/bin/env python\n'
u"#!/usr/bin/python2\n",
u"#! /usr/bin/python3\n",
]
not_shebang_comments = [u"# I saw a giant python\n",
u"# I have never seen a python2\n",
]
for comment in shebang_comments:
node = FromImport(u'math', [Leaf(token.NAME, u'cos', prefix=" ")])
node.prefix = comment
self.assertTrue(is_shebang_comment(node))
for comment in not_shebang_comments:
node = FromImport(u'math', [Leaf(token.NAME, u'cos', prefix=" ")])
node.prefix = comment
self.assertFalse(is_shebang_comment(node))
def test_is_encoding_comment(self):
"""
Tests whether the fixer_util.is_encoding_comment() function is working.
"""
encoding_comments = [u"# coding: utf-8",
u"# encoding: utf-8",
u"# -*- coding: latin-1 -*-",
u"# vim: set fileencoding=iso-8859-15 :",
]
not_encoding_comments = [u"# We use the file encoding utf-8",
u"coding = 'utf-8'",
u"encoding = 'utf-8'",
]
for comment in encoding_comments:
node = FromImport(u'math', [Leaf(token.NAME, u'cos', prefix=" ")])
node.prefix = comment
self.assertTrue(is_encoding_comment(node))
for comment in not_encoding_comments:
node = FromImport(u'math', [Leaf(token.NAME, u'cos', prefix=" ")])
node.prefix = comment
self.assertFalse(is_encoding_comment(node))
def touch_import_top(package, name_to_import, node):
"""Works like `does_tree_import` but adds an import statement at the
top if it was not imported (but below any __future__ imports) and below any
comments such as shebang lines).
Based on lib2to3.fixer_util.touch_import()
Calling this multiple times adds the imports in reverse order.
Also adds "standard_library.install_aliases()" after "from future import
standard_library". This should probably be factored into another function.
"""
root = find_root(node)
if does_tree_import(package, name_to_import, root):
return
# Ideally, we would look for whether futurize --all-imports has been run,
# as indicated by the presence of ``from builtins import (ascii, ...,
# zip)`` -- and, if it has, we wouldn't import the name again.
# Look for __future__ imports and insert below them
found = False
for name in [
'absolute_import', 'division', 'print_function', 'unicode_literals'
]:
if does_tree_import('__future__', name, root):
found = True
break
if found:
# At least one __future__ import. We want to loop until we've seen them
# all.
start, end = None, None
for idx, node in enumerate(root.children):
if check_future_import(node):
start = idx
# Start looping
idx2 = start
while node:
node = node.next_sibling
idx2 += 1
if not check_future_import(node):
end = idx2
break
break
assert start is not None
assert end is not None
insert_pos = end
else:
# No __future__ imports.
# We look for a docstring and insert the new node below that. If no docstring
# exists, just insert the node at the top.
for idx, node in enumerate(root.children):
if node.type != syms.simple_stmt:
break
if not is_docstring(node):
# This is the usual case.
break
insert_pos = idx
if package is None:
import_ = Node(syms.import_name, [
Leaf(token.NAME, u"import"),
Leaf(token.NAME, name_to_import, prefix=u" ")
])
else:
import_ = FromImport(package,
[Leaf(token.NAME, name_to_import, prefix=u" ")])
if name_to_import == u'standard_library':
# Add:
# standard_library.install_aliases()
# after:
# from future import standard_library
install_hooks = Node(syms.simple_stmt, [
Node(syms.power, [
Leaf(token.NAME, u'standard_library'),
Node(syms.trailer, [
Leaf(token.DOT, u'.'),
Leaf(token.NAME, u'install_aliases')
]),
Node(syms.trailer,
[Leaf(token.LPAR, u'('),
Leaf(token.RPAR, u')')])
])
])
children_hooks = [install_hooks, Newline()]
else:
children_hooks = []
# FromImport(package, [Leaf(token.NAME, name_to_import, prefix=u" ")])
children_import = [import_, Newline()]
old_prefix = root.children[insert_pos].prefix
root.children[insert_pos].prefix = u''
root.insert_child(
insert_pos, Node(syms.simple_stmt, children_import, prefix=old_prefix))
if len(children_hooks) > 0:
root.insert_child(insert_pos + 1, Node(syms.simple_stmt,
children_hooks))
def transform_member(self, node, results):
"""Transform for imports of specific module elements. Replaces
the module to be imported from with the appropriate new
module.
"""
mod_member = results.get("mod_member")
pref = mod_member.prefix
member = results.get("member")
# Simple case with only a single member being imported
if member:
# this may be a list of length one, or just a node
if isinstance(member, list):
member = member[0]
new_name = None
for change in MAPPING[mod_member.value]:
if member.value in change[1]:
new_name = change[0]
break
if new_name:
mod_member.replace(Name(new_name, prefix=pref))
else:
self.cannot_convert(node, "This is an invalid module element")
# Multiple members being imported
else:
# a dictionary for replacements, order matters
modules = []
mod_dict = {}
members = results["members"]
for member in members:
# we only care about the actual members
if member.type == syms.import_as_name:
as_name = member.children[2].value
member_name = member.children[0].value
else:
member_name = member.value
as_name = None
if member_name != u",":
for change in MAPPING[mod_member.value]:
if member_name in change[1]:
if change[0] not in mod_dict:
modules.append(change[0])
mod_dict.setdefault(change[0], []).append(member)
new_nodes = []
indentation = find_indentation(node)
first = True
def handle_name(name, prefix):
if name.type == syms.import_as_name:
kids = [
Name(name.children[0].value, prefix=prefix),
name.children[1].clone(), name.children[2].clone()
]
return [Node(syms.import_as_name, kids)]
return [Name(name.value, prefix=prefix)]
for module in modules:
elts = mod_dict[module]
names = []
for elt in elts[:-1]:
names.extend(handle_name(elt, pref))
names.append(Comma())
names.extend(handle_name(elts[-1], pref))
new = FromImport(module, names)
if not first or node.parent.prefix.endswith(indentation):
new.prefix = indentation
new_nodes.append(new)
first = False
if new_nodes:
nodes = []
for new_node in new_nodes[:-1]:
nodes.extend([new_node, Newline()])
nodes.append(new_nodes[-1])
node.replace(nodes)
else:
self.cannot_convert(node, "All module elements are invalid")
def transform(self, node, results):
# The patterns dictate which of these names will be defined
name = results.get("name")
attr = results.get("attr")
if attr is None:
attr = Name("__init__")
using = results.get("using")
in_list = results.get("in_list")
imp_list = results.get("imp_list")
power = results.get("pow")
before = results.get("before")
after = results.get("after")
d_name = results.get("dotted_name")
# An import_stmt is always contained within a simple_stmt
simple_stmt = node.parent
# The parent is useful for adding new import_stmts
parent = simple_stmt.parent
idx = parent.children.index(simple_stmt)
if any((results.get("from_import_rename")
is not None, results.get("name_import_rename") is not None)):
self.cannot_convert(node,
reason="ambiguity: import binds a single name")
elif using is None and not in_list:
# import urllib.request, single-name import
replacement = name_import_replacement(name, attr)
replacement.prefix = node.prefix
node.replace(replacement)
elif using is None:
# import ..., urllib.request, math, http.sever, ...
for d_name in imp_list:
if d_name.type == syms.dotted_name:
name = d_name.children[0]
attr = d_name.children[2]
elif d_name.type == token.NAME and d_name.value + ".__init__" in MAPPING:
name = d_name
attr = Name("__init__")
else:
continue
if name.value + "." + attr.value not in MAPPING:
continue
candidates = all_candidates(name.value, attr.value)
children = [Name("import")]
for c in candidates:
children.append(Name(c, prefix=" "))
children.append(Comma())
children.pop()
# Put in the new statement.
indent = indentation(simple_stmt)
next_stmt = Node(syms.simple_stmt,
[Node(syms.import_name, children),
Newline()])
parent.insert_child(idx + 1, next_stmt)
parent.insert_child(idx + 1, Leaf(token.INDENT, indent))
# Remove the old imported name
test_comma = d_name.next_sibling
if test_comma and test_comma.type == token.COMMA:
test_comma.remove()
elif test_comma is None:
test_comma = d_name.prev_sibling
if test_comma and test_comma.type == token.COMMA:
test_comma.remove()
d_name.remove()
if not in_list.children:
simple_stmt.remove()
elif in_list is not None:
##########################################################
# "from urllib.request import urlopen, urlretrieve, ..." #
# Replace one import statement with potentially many. #
##########################################################
packages = dict([(n, [])
for n in all_candidates(name.value, attr.value)])
# Figure out what names need to be imported from what
# Add them to a dict to be parsed once we're completely done
for imported in using:
if imported.type == token.COMMA:
continue
if imported.type == syms.import_as_name:
test_name = imported.children[0].value
if len(imported.children) > 2:
# 'as' whatever
rename = imported.children[2].value
else:
rename = None
elif imported.type == token.NAME:
test_name = imported.value
rename = None
pkg = new_package(name.value, attr.value, test_name)
packages[pkg].append((test_name, rename))
# Parse the dict to create new import statements to replace this one
imports = []
for new_pkg, names in packages.items():
if not names:
# Didn't import anything from that package, move along
continue
new_names = []
for test_name, rename in names:
if rename is None:
new_names.append(Name(test_name, prefix=" "))
else:
new_names.append(
ImportAsName(test_name, rename, prefix=" "))
new_names.append(Comma())
new_names.pop()
imports.append(FromImport(new_pkg, new_names))
# Replace this import statement with one of the others
replacement = imports.pop()
replacement.prefix = node.prefix
node.replace(replacement)
indent = indentation(simple_stmt)
# Add the remainder of the imports as new statements.
while imports:
next_stmt = Node(syms.simple_stmt, [imports.pop(), Newline()])
parent.insert_child(idx + 1, next_stmt)
parent.insert_child(idx + 1, Leaf(token.INDENT, indent))
elif using.type == token.STAR:
# from urllib.request import *
nodes = [
FromImport(pkg, [Star(prefix=" ")])
for pkg in all_candidates(name.value, attr.value)
]
replacement = nodes.pop()
replacement.prefix = node.prefix
node.replace(replacement)
indent = indentation(simple_stmt)
while nodes:
next_stmt = Node(syms.simple_stmt, [nodes.pop(), Newline()])
parent.insert_child(idx + 1, next_stmt)
parent.insert_child(idx + 1, Leaf(token.INDENT, indent))
elif power is not None:
# urllib.request.urlopen
# Replace it with urllib2.urlopen
pkg = new_package(name.value, attr.value, using.value)
# Remove the trailer node that contains attr.
if pkg:
if attr.parent:
attr.parent.remove()
name.replace(Name(pkg, prefix=name.prefix))
elif using.type == token.NAME:
# from urllib.request import urlopen
pkg = new_package(name.value, attr.value, using.value)
if attr.value == "__init__" and pkg == name.value:
# Replacing "from abc import xyz" with "from abc import xyz"
# Just leave it alone so as not to mess with other fixers
return
else:
node.replace(FromImport(pkg, [using]))
def transform_member(self, node, results):
"""Transform for imports of specific module elements. Replaces
the module to be imported from with the appropriate new
module.
"""
mod_member = results.get('mod_member')
pref = mod_member.prefix
member = results.get('member')
if member:
if isinstance(member, list):
member = member[0]
new_name = None
for change in MAPPING[mod_member.value]:
if member.value in change[1]:
new_name = change[0]
break
if new_name:
mod_member.replace(Name(new_name, prefix=pref))
else:
self.cannot_convert(node, 'This is an invalid module element')
else:
modules = []
mod_dict = {}
members = results['members']
for member in members:
if member.type == syms.import_as_name:
as_name = member.children[2].value
member_name = member.children[0].value
else:
member_name = member.value
as_name = None
if member_name != u',':
for change in MAPPING[mod_member.value]:
if member_name in change[1]:
if change[0] not in mod_dict:
modules.append(change[0])
mod_dict.setdefault(change[0], []).append(member)
new_nodes = []
indentation = find_indentation(node)
first = True
def handle_name(name, prefix):
if name.type == syms.import_as_name:
kids = [Name(name.children[0].value, prefix=prefix), name.children[1].clone(), name.children[2].clone()]
return [Node(syms.import_as_name, kids)]
return [Name(name.value, prefix=prefix)]
for module in modules:
elts = mod_dict[module]
names = []
for elt in elts[:-1]:
names.extend(handle_name(elt, pref))
names.append(Comma())
names.extend(handle_name(elts[-1], pref))
new = FromImport(module, names)
if not first or node.parent.prefix.endswith(indentation):
new.prefix = indentation
new_nodes.append(new)
first = False
if new_nodes:
nodes = []
for new_node in new_nodes[:-1]:
nodes.extend([new_node, Newline()])
nodes.append(new_nodes[-1])
node.replace(nodes)
else:
self.cannot_convert(node, 'All module elements are invalid')
return
def transform_member(self, node, results):
"""Transform for imports of specific module elements. Replaces
the module to be imported from with the appropriate new
module.
"""
mod_member = results.get("mod_member")
pref = mod_member.prefix
member = results.get("member")
# Simple case with only a single member being imported
if member:
# this may be a list of length one, or just a node
if isinstance(member, list):
member = member[0]
new_name = None
for change in MAPPING[mod_member.value]:
if member.value in change[1]:
new_name = change[0]
break
if new_name:
mod_member.replace(Name(new_name, prefix=pref))
else:
self.cannot_convert(node, "This is an invalid module element")
# Multiple members being imported
else:
# a dictionary for replacements, order matters
modules = []
mod_dict = {}
members = results["members"]
for member in members:
# we only care about the actual members
if member.type == syms.import_as_name:
as_name = member.children[2].value
member_name = member.children[0].value
else:
member_name = member.value
as_name = None
if member_name != ",":
for change in MAPPING[mod_member.value]:
if member_name in change[1]:
if change[0] not in mod_dict:
modules.append(change[0])
mod_dict.setdefault(change[0], []).append(member)
new_nodes = []
indentation = find_indentation(node)
first = True
def handle_name(name, prefix):
if name.type == syms.import_as_name:
kids = [Name(name.children[0].value, prefix=prefix),
name.children[1].clone(),
name.children[2].clone()]
return [Node(syms.import_as_name, kids)]
return [Name(name.value, prefix=prefix)]
for module in modules:
elts = mod_dict[module]
names = []
for elt in elts[:-1]:
names.extend(handle_name(elt, pref))
names.append(Comma())
names.extend(handle_name(elts[-1], pref))
new = FromImport(module, names)
if not first or node.parent.prefix.endswith(indentation):
new.prefix = indentation
new_nodes.append(new)
first = False
if new_nodes:
nodes = []
for new_node in new_nodes[:-1]:
nodes.extend([new_node, Newline()])
nodes.append(new_nodes[-1])
node.replace(nodes)
else:
self.cannot_convert(node, "All module elements are invalid")
