EVOLUTION-MANAGER

Edit File: import_utils.py

# coding=utf-8 """Functions for dealing with import statements.""" # Copyright 2017 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from __future__ import absolute_import from __future__ import division from __future__ import print_function import ast import copy import logging from pasta.augment import errors from pasta.base import ast_utils from pasta.base import scope def add_import(tree, name_to_import, asname=None, from_import=True, merge_from_imports=True): """Adds an import to the module. This function will try to ensure not to create duplicate imports. If name_to_import is already imported, it will return the existing import. This is true even if asname is set (asname will be ignored, and the existing name will be returned). If the import would create a name that already exists in the scope given by tree, this function will "import as", and append "_x" to the asname where x is the smallest positive integer generating a unique name. Arguments: tree: (ast.Module) Module AST to modify. name_to_import: (string) The absolute name to import. asname: (string) The alias for the import ("import name_to_import as asname") from_import: (boolean) If True, import the name using an ImportFrom node. merge_from_imports: (boolean) If True, merge a newly inserted ImportFrom node into an existing ImportFrom node, if applicable. Returns: The name (as a string) that can be used to reference the imported name. This can be the fully-qualified name, the basename, or an alias name. """ sc = scope.analyze(tree) # Don't add anything if it's already imported if name_to_import in sc.external_references: existing_ref = next((ref for ref in sc.external_references[name_to_import] if ref.name_ref is not None), None) if existing_ref: return existing_ref.name_ref.id import_node = None added_name = None def make_safe_alias_node(alias_name, asname): # Try to avoid name conflicts new_alias = ast.alias(name=alias_name, asname=asname) imported_name = asname or alias_name counter = 0 while imported_name in sc.names: counter += 1 imported_name = new_alias.asname = '%s_%d' % (asname or alias_name, counter) return new_alias # Add an ImportFrom node if requested and possible if from_import and '.' in name_to_import: from_module, alias_name = name_to_import.rsplit('.', 1) new_alias = make_safe_alias_node(alias_name, asname) if merge_from_imports: # Try to add to an existing ImportFrom from the same module existing_from_import = next( (node for node in tree.body if isinstance(node, ast.ImportFrom) and node.module == from_module and node.level == 0), None) if existing_from_import: existing_from_import.names.append(new_alias) return new_alias.asname or new_alias.name # Create a new node for this import import_node = ast.ImportFrom(module=from_module, names=[new_alias], level=0) # If not already created as an ImportFrom, create a normal Import node if not import_node: new_alias = make_safe_alias_node(alias_name=name_to_import, asname=asname) import_node = ast.Import( names=[new_alias]) # Insert the node at the top of the module and return the name in scope tree.body.insert(1 if ast_utils.has_docstring(tree) else 0, import_node) return new_alias.asname or new_alias.name def split_import(sc, node, alias_to_remove): """Split an import node by moving the given imported alias into a new import. Arguments: sc: (scope.Scope) Scope computed on whole tree of the code being modified. node: (ast.Import|ast.ImportFrom) An import node to split. alias_to_remove: (ast.alias) The import alias node to remove. This must be a child of the given `node` argument. Raises: errors.InvalidAstError: if `node` is not appropriately contained in the tree represented by the scope `sc`. """ parent = sc.parent(node) parent_list = None for a in ('body', 'orelse', 'finalbody'): if hasattr(parent, a) and node in getattr(parent, a): parent_list = getattr(parent, a) break else: raise errors.InvalidAstError('Unable to find list containing import %r on ' 'parent node %r' % (node, parent)) idx = parent_list.index(node) new_import = copy.deepcopy(node) new_import.names = [alias_to_remove] node.names.remove(alias_to_remove) parent_list.insert(idx + 1, new_import) return new_import def get_unused_import_aliases(tree, sc=None): """Get the import aliases that aren't used. Arguments: tree: (ast.AST) An ast to find imports in. sc: A scope.Scope representing tree (generated from scratch if not provided). Returns: A list of ast.alias representing imported aliases that aren't referenced in the given tree. """ if sc is None: sc = scope.analyze(tree) unused_aliases = set() for node in ast.walk(tree): if isinstance(node, ast.alias): str_name = node.asname if node.asname is not None else node.name if str_name in sc.names: name = sc.names[str_name] if not name.reads: unused_aliases.add(node) else: # This happens because of https://github.com/google/pasta/issues/32 logging.warning('Imported name %s not found in scope (perhaps it\'s ' 'imported dynamically)', str_name) return unused_aliases def remove_import_alias_node(sc, node): """Remove an alias and if applicable remove their entire import. Arguments: sc: (scope.Scope) Scope computed on whole tree of the code being modified. node: (ast.Import|ast.ImportFrom|ast.alias) The node to remove. """ import_node = sc.parent(node) if len(import_node.names) == 1: import_parent = sc.parent(import_node) ast_utils.remove_child(import_parent, import_node) else: ast_utils.remove_child(import_node, node) def remove_duplicates(tree, sc=None): """Remove duplicate imports, where it is safe to do so. This does NOT remove imports that create new aliases Arguments: tree: (ast.AST) An ast to modify imports in. sc: A scope.Scope representing tree (generated from scratch if not provided). Returns: Whether any changes were made. """ if sc is None: sc = scope.analyze(tree) modified = False seen_names = set() for node in tree.body: if isinstance(node, (ast.Import, ast.ImportFrom)): for alias in list(node.names): import_node = sc.parent(alias) if isinstance(import_node, ast.Import): full_name = alias.name elif import_node.module: full_name = '%s%s.%s' % ('.' * import_node.level, import_node.module, alias.name) else: full_name = '%s%s' % ('.' * import_node.level, alias.name) full_name += ':' + (alias.asname or alias.name) if full_name in seen_names: remove_import_alias_node(sc, alias) modified = True else: seen_names.add(full_name) return modified