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 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 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 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 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 transform(self, node, results):
assert results
# Not sure if this check is actually needed, and having it means
# we can't remove these imports. Oh well.
if not does_tree_import('__future__', 'print_function', node):
return
assert node.children[0] == Name(u"print")
args = node.children[1:]
if len(args) != 1:
raise RuntimeError('I didn\'t expect this')
args_node = args[0].clone()
if args_node:
if args_node.children[0] == Leaf(token.LPAR, '(') \
and args_node.children[-1] == Leaf(token.RPAR, ')'):
args_node.children[0].remove()
args_node.children[-1].remove()
args_node.prefix = u""
log_call_args = args_node.children
if len(log_call_args) == 1 and log_call_args[0].type == syms.arglist:
log_call_args = log_call_args[0].children
if any(node.type == syms.argument for node in log_call_args):
raise RuntimeError(
'Why are there kwargs in here!? How do we handle them? %s',
log_call_args)
return self._create_logging_call(log_call_args, node)
def add_globals(self, tree, inserted_types):
"""Add required globals to the tree. Idempotent."""
# Copy imports if not already present
top_lines = []
def import_name(name, alias):
return name + ('' if alias is None else ' as %s' % alias)
for (pkg, pkg_alias), names in self._get_imports(inserted_types):
if not names:
if does_tree_import(None, pkg_alias or pkg, tree):
continue
top_lines.append('import %s\n' % import_name(pkg, pkg_alias))
else:
assert pkg_alias is None
import_names = []
for name, alias in names:
if does_tree_import(pkg, alias or name, tree):
continue
import_names.append(import_name(name, alias))
if not import_names:
continue
top_lines.append('from %s import %s\n' %
(pkg, ', '.join(import_names)))
import_idx = [
idx for idx, idx_node in enumerate(tree.children)
if self.import_pattern.match(idx_node)
]
if import_idx:
insert_pos = import_idx[-1] + 1
else:
insert_pos = 0
# first string (normally docstring)
for idx, idx_node in enumerate(tree.children):
if (idx_node.type == syms.simple_stmt and idx_node.children
and idx_node.children[0].type == token.STRING):
insert_pos = idx + 1
break
if self.assignments:
top_lines.append('\n')
top_lines.extend(str(a).strip() + '\n' for a in self.assignments)
top_lines = Util.parse_string(''.join(top_lines))
for offset, offset_node in enumerate(top_lines.children[:-1]):
tree.insert_child(insert_pos + offset, offset_node)
def type_updater(self, match, node):
# type: (Match, Node) -> str
# Replace `pkg.mod.SomeClass` with `SomeClass`
# and remember to import it.
word = match.group()
if does_tree_import(None, word, node):
# Check whether there already exists an import binding for this
return word
# If not, assume it's either builtin or from `typing`
self.touch_typing_import(word, node)
return word
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 add_import(import_name, node):
suite = get_parent_of_type(node, syms.suite)
test_case = suite
while test_case.parent.type != syms.file_input:
test_case = test_case.parent
file_input = test_case.parent
if not does_tree_import(None, import_name, node):
import_stmt = Node(syms.simple_stmt,
[Node(syms.import_name, [Name('import'), Name(import_name, prefix=' ')]),
Newline(),
])
insert_import(import_stmt, test_case, file_input)
def add_import(import_name, node):
suite = get_parent_of_type(node, syms.suite)
test_case = suite
while test_case.parent.type != syms.file_input:
test_case = test_case.parent
file_input = test_case.parent
if not does_tree_import(None, import_name, node):
import_stmt = Node(syms.simple_stmt, [
Node(
syms.import_name,
[Name('import'), Name(import_name, prefix=' ')]),
Newline(),
])
insert_import(import_stmt, test_case, file_input)
def transform(self, node, results):
from_import = results.get(u"from_import")
from_import_submod = results.get(u"from_import_submod")
name_import = results.get(u"name_import")
dotted_name = results.get(u"dotted_name")
name = results.get(u"name")
names = results.get(u"names")
attr = results.get(u"attr")
imported = results.get(u"imported")
if names:
for name in names:
if name.type == token.NAME:
self.fix_simple_name(name)
elif name.type == syms.dotted_as_name:
self.fix_simple_name(name.children[0]) if name.children[0].type == token.NAME else \
self.fix_dotted_name(name.children[0])
elif name.type == syms.dotted_name:
self.fix_dotted_name(name)
elif from_import_submod:
renamed = results.get(u"renamed")
new_name, new_attr = self.get_dotted_import_replacement(
name, attr, renamed=renamed)
if new_attr is not None:
name.replace(new_name)
attr.replace(new_attr)
else:
children = [Name(u"import"), new_name]
node.replace(
Node(syms.import_name, children, prefix=node.prefix))
elif dotted_name:
self.fix_dotted_name(dotted_name)
elif name_import or from_import:
self.fix_simple_name(name)
elif name and not attr:
if does_tree_import(None, MAPPING[name.value], node):
self.fix_simple_name(name)
elif name and attr:
# Note that this will fix a dotted name that was never imported.
# This will probably not matter.
self.fix_dotted_name(node)
elif imported and imported.type == syms.import_as_names:
self.fix_submod_import(imported=imported.children,
node=node,
name=name.value)
def transform(self, node, results):
from_import = results.get("from_import")
from_import_submod = results.get("from_import_submod")
name_import = results.get("name_import")
dotted_name = results.get("dotted_name")
name = results.get("name")
names = results.get("names")
attr = results.get("attr")
imported = results.get("imported")
if names:
for name in names:
if name.type == token.NAME:
self.fix_simple_name(name)
elif name.type == syms.dotted_as_name:
self.fix_simple_name(name.children[0]) if name.children[0].type == token.NAME else \
self.fix_dotted_name(name.children[0])
elif name.type == syms.dotted_name:
self.fix_dotted_name(name)
elif from_import_submod:
renamed = results.get("renamed")
new_name, new_attr = self.get_dotted_import_replacement(name, attr, renamed=renamed)
if new_attr is not None:
name.replace(new_name)
attr.replace(new_attr)
else:
children = [Name("import"), new_name]
node.replace(Node(syms.import_name, children, prefix=node.prefix))
elif dotted_name:
self.fix_dotted_name(dotted_name)
elif name_import or from_import:
self.fix_simple_name(name)
elif name and not attr:
if does_tree_import(None, MAPPING[name.value], node) and \
is_probably_builtin(name):
self.fix_simple_name(name)
elif name and attr:
# Note that this will fix a dotted name that was never imported. This will probably not matter.
self.fix_dotted_name(node)
elif imported and imported.type == syms.import_as_names:
self.fix_submod_import(imported=imported.children, node=node, name=name.value)
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 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 does_tree_import(self, package, name, string):
node = parse(string)
node = self._find_bind_rec('start', node)
return fixer_util.does_tree_import(package, name, 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).
Based on lib2to3.fixer_util.touch_import()
Calling this multiple times adds the imports in reverse order.
Also adds "standard_library.install_hooks()" 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 future.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 (node.children and node.children[0].type == token.STRING):
# 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_hooks()
# 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_hooks')]),
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()]
root.insert_child(insert_pos, Node(syms.simple_stmt, children_import))
if len(children_hooks) > 0:
root.insert_child(insert_pos + 1, Node(syms.simple_stmt, children_hooks))
def transform(self, node, results):
if does_tree_import('builtins', 'input', node):
return
return super(FixInput, self).transform(node, results)
def does_tree_import(self, package, name, string):
node = parse(string)
# Find the binding of start -- that's what we'll go from
node = self._find_bind_rec('start', node)
return fixer_util.does_tree_import(package, name, node)
def does_tree_import(self, package, name, string):
node = parse(string)
# Find the binding of start -- that's what we'll go from
node = self._find_bind_rec("start", node)
return fixer_util.does_tree_import(package, name, node)
