Source code for libcst.codemod.visitors._add_imports

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from collections import defaultdict
from typing import Dict, List, Optional, Sequence, Set, Tuple, Union

import libcst
from libcst import matchers as m, parse_statement
from libcst._nodes.statement import Import, ImportFrom, SimpleStatementLine
from libcst.codemod._context import CodemodContext
from libcst.codemod._visitor import ContextAwareTransformer
from libcst.codemod.visitors._gather_imports import _GatherImportsMixin
from libcst.codemod.visitors._imports import ImportItem
from libcst.helpers import get_absolute_module_from_package_for_import
from libcst.helpers.common import ensure_type


class _GatherTopImportsBeforeStatements(_GatherImportsMixin):
    """
    Works similarly to GatherImportsVisitor, but only considers imports
    declared before any other statements of the module with the exception
    of docstrings and __strict__ flag.
    """

    def __init__(self, context: CodemodContext) -> None:
        super().__init__(context)
        # Track all of the imports found in this transform
        self.all_imports: List[Union[libcst.Import, libcst.ImportFrom]] = []

    def leave_Module(self, original_node: libcst.Module) -> None:
        start = 1 if _skip_first(original_node) else 0
        for stmt in original_node.body[start:]:
            if m.matches(
                stmt,
                m.SimpleStatementLine(body=[m.ImportFrom() | m.Import()]),
            ):
                stmt = ensure_type(stmt, SimpleStatementLine)
                # Workaround for python 3.8 and 3.9, won't accept Union for isinstance
                if m.matches(stmt.body[0], m.ImportFrom()):
                    imp = ensure_type(stmt.body[0], ImportFrom)
                    self.all_imports.append(imp)
                if m.matches(stmt.body[0], m.Import()):
                    imp = ensure_type(stmt.body[0], Import)
                    self.all_imports.append(imp)
            else:
                break
        for imp in self.all_imports:
            if m.matches(imp, m.Import()):
                imp = ensure_type(imp, Import)
                self._handle_Import(imp)
            else:
                imp = ensure_type(imp, ImportFrom)
                self._handle_ImportFrom(imp)


[docs] class AddImportsVisitor(ContextAwareTransformer): """ Ensures that given imports exist in a module. Given a :class:`~libcst.codemod.CodemodContext` and a sequence of tuples specifying a module to import from as a string. Optionally an object to import from that module and any alias to assign that import, ensures that import exists. It will modify existing imports as necessary if the module in question is already being imported from. This is one of the transforms that is available automatically to you when running a codemod. To use it in this manner, import :class:`~libcst.codemod.visitors.AddImportsVisitor` and then call the static :meth:`~libcst.codemod.visitors.AddImportsVisitor.add_needed_import` method, giving it the current context (found as ``self.context`` for all subclasses of :class:`~libcst.codemod.Codemod`), the module you wish to import from and optionally an object you wish to import from that module and any alias you would like to assign that import to. For example:: AddImportsVisitor.add_needed_import(self.context, "typing", "Optional") This will produce the following code in a module, assuming there was no typing import already:: from typing import Optional As another example:: AddImportsVisitor.add_needed_import(self.context, "typing") This will produce the following code in a module, assuming there was no import already:: import typing Note that this is a subclass of :class:`~libcst.CSTTransformer` so it is possible to instantiate it and pass it to a :class:`~libcst.Module` :meth:`~libcst.CSTNode.visit` method. However, it is far easier to use the automatic transform feature of :class:`~libcst.codemod.CodemodCommand` and schedule an import to be added by calling :meth:`~libcst.codemod.visitors.AddImportsVisitor.add_needed_import` """ CONTEXT_KEY = "AddImportsVisitor" @staticmethod def _get_imports_from_context( context: CodemodContext, ) -> List[ImportItem]: imports = context.scratch.get(AddImportsVisitor.CONTEXT_KEY, []) if not isinstance(imports, list): raise Exception("Logic error!") return imports
[docs] @staticmethod def add_needed_import( context: CodemodContext, module: str, obj: Optional[str] = None, asname: Optional[str] = None, relative: int = 0, ) -> None: """ Schedule an import to be added in a future invocation of this class by updating the ``context`` to include the ``module`` and optionally ``obj`` to be imported as well as optionally ``alias`` to alias the imported ``module`` or ``obj`` to. When subclassing from :class:`~libcst.codemod.CodemodCommand`, this will be performed for you after your transform finishes executing. If you are subclassing from a :class:`~libcst.codemod.Codemod` instead, you will need to call the :meth:`~libcst.codemod.Codemod.transform_module` method on the module under modification with an instance of this class after performing your transform. Note that if the particular ``module`` or ``obj`` you are requesting to import already exists as an import on the current module at the time of executing :meth:`~libcst.codemod.Codemod.transform_module` on an instance of :class:`~libcst.codemod.visitors.AddImportsVisitor`, this will perform no action in order to avoid adding duplicate imports. """ if module == "__future__" and obj is None: raise Exception("Cannot import __future__ directly!") imports = AddImportsVisitor._get_imports_from_context(context) imports.append(ImportItem(module, obj, asname, relative)) context.scratch[AddImportsVisitor.CONTEXT_KEY] = imports
def __init__( self, context: CodemodContext, imports: Sequence[ImportItem] = (), ) -> None: # Allow for instantiation from either a context (used when multiple transforms # get chained) or from a direct instantiation. super().__init__(context) imps: List[ImportItem] = [ *AddImportsVisitor._get_imports_from_context(context), *imports, ] # Verify that the imports are valid for imp in imps: if imp.module == "__future__" and imp.obj_name is None: raise Exception("Cannot import __future__ directly!") if imp.module == "__future__" and imp.alias is not None: raise Exception("Cannot import __future__ objects with aliases!") # Resolve relative imports if we have a module name imps = [imp.resolve_relative(self.context.full_package_name) for imp in imps] # List of modules we need to ensure are imported self.module_imports: Set[str] = { imp.module for imp in imps if imp.obj_name is None and imp.alias is None } # List of modules we need to check for object imports on from_imports: Set[str] = { imp.module for imp in imps if imp.obj_name is not None and imp.alias is None } # Mapping of modules we're adding to the object they should import self.module_mapping: Dict[str, Set[str]] = { module: { imp.obj_name for imp in imps if imp.module == module and imp.obj_name is not None and imp.alias is None } for module in sorted(from_imports) } # List of aliased modules we need to ensure are imported self.module_aliases: Dict[str, str] = { imp.module: imp.alias for imp in imps if imp.obj_name is None and imp.alias is not None } # List of modules we need to check for object imports on from_imports_aliases: Set[str] = { imp.module for imp in imps if imp.obj_name is not None and imp.alias is not None } # Mapping of modules we're adding to the object with alias they should import self.alias_mapping: Dict[str, List[Tuple[str, str]]] = { module: [ (imp.obj_name, imp.alias) for imp in imps if imp.module == module and imp.obj_name is not None and imp.alias is not None ] for module in sorted(from_imports_aliases) } # Track the list of imports found at the top of the file self.all_imports: List[Union[libcst.Import, libcst.ImportFrom]] = [] def visit_Module(self, node: libcst.Module) -> None: # Do a preliminary pass to gather the imports we already have at the top gatherer = _GatherTopImportsBeforeStatements(self.context) node.visit(gatherer) self.all_imports = gatherer.all_imports self.module_imports = self.module_imports - gatherer.module_imports for module, alias in gatherer.module_aliases.items(): if module in self.module_aliases and self.module_aliases[module] == alias: del self.module_aliases[module] for module, aliases in gatherer.alias_mapping.items(): for obj, alias in aliases: if ( module in self.alias_mapping and (obj, alias) in self.alias_mapping[module] ): self.alias_mapping[module].remove((obj, alias)) if len(self.alias_mapping[module]) == 0: del self.alias_mapping[module] for module, imports in gatherer.object_mapping.items(): if module not in self.module_mapping: # We don't care about this import at all continue elif "*" in imports: # We already implicitly are importing everything del self.module_mapping[module] else: # Lets figure out what's left to import self.module_mapping[module] = self.module_mapping[module] - imports if not self.module_mapping[module]: # There's nothing left, so lets delete this work item del self.module_mapping[module] def leave_ImportFrom( self, original_node: libcst.ImportFrom, updated_node: libcst.ImportFrom ) -> libcst.ImportFrom: if isinstance(updated_node.names, libcst.ImportStar): # There's nothing to do here! return updated_node # Ensure this is one of the imports at the top if original_node not in self.all_imports: return updated_node # Get the module we're importing as a string, see if we have work to do. module = get_absolute_module_from_package_for_import( self.context.full_package_name, updated_node ) if ( module is None or module not in self.module_mapping and module not in self.alias_mapping ): return updated_node # We have work to do, mark that we won't modify this again. imports_to_add = self.module_mapping.get(module, []) if module in self.module_mapping: del self.module_mapping[module] aliases_to_add = self.alias_mapping.get(module, []) if module in self.alias_mapping: del self.alias_mapping[module] # Now, do the actual update. return updated_node.with_changes( names=[ *( libcst.ImportAlias(name=libcst.Name(imp)) for imp in sorted(imports_to_add) ), *( libcst.ImportAlias( name=libcst.Name(imp), asname=libcst.AsName(name=libcst.Name(alias)), ) for (imp, alias) in sorted(aliases_to_add) ), *updated_node.names, ] ) def _split_module( self, orig_module: libcst.Module, updated_module: libcst.Module ) -> Tuple[ List[Union[libcst.SimpleStatementLine, libcst.BaseCompoundStatement]], List[Union[libcst.SimpleStatementLine, libcst.BaseCompoundStatement]], List[Union[libcst.SimpleStatementLine, libcst.BaseCompoundStatement]], ]: statement_before_import_location = 0 import_add_location = 0 # This works under the principle that while we might modify node contents, # we have yet to modify the number of statements. So we can match on the # original tree but break up the statements of the modified tree. If we # change this assumption in this visitor, we will have to change this code. # Finds the location to add imports. It is the end of the first import block that occurs before any other statement (save for docstrings) # Never insert an import before initial __strict__ flag or docstring if _skip_first(orig_module): statement_before_import_location = import_add_location = 1 for i, statement in enumerate( orig_module.body[statement_before_import_location:] ): if m.matches( statement, m.SimpleStatementLine(body=[m.ImportFrom() | m.Import()]) ): import_add_location = i + statement_before_import_location + 1 else: break return ( list(updated_module.body[:statement_before_import_location]), list( updated_module.body[ statement_before_import_location:import_add_location ] ), list(updated_module.body[import_add_location:]), ) def _insert_empty_line( self, statements: List[ Union[libcst.SimpleStatementLine, libcst.BaseCompoundStatement] ], ) -> List[Union[libcst.SimpleStatementLine, libcst.BaseCompoundStatement]]: if len(statements) < 1: # No statements, nothing to add to return statements if len(statements[0].leading_lines) == 0: # Statement has no leading lines, add one! return [ statements[0].with_changes(leading_lines=(libcst.EmptyLine(),)), *statements[1:], ] if statements[0].leading_lines[0].comment is None: # First line is empty, so its safe to leave as-is return statements # Statement has a comment first line, so lets add one more empty line return [ statements[0].with_changes( leading_lines=(libcst.EmptyLine(), *statements[0].leading_lines) ), *statements[1:], ] def leave_Module( self, original_node: libcst.Module, updated_node: libcst.Module ) -> libcst.Module: # Don't try to modify if we have nothing to do if ( not self.module_imports and not self.module_mapping and not self.module_aliases and not self.alias_mapping ): return updated_node # First, find the insertion point for imports ( statements_before_imports, statements_until_add_imports, statements_after_imports, ) = self._split_module(original_node, updated_node) # Make sure there's at least one empty line before the first non-import statements_after_imports = self._insert_empty_line(statements_after_imports) # Mapping of modules we're adding to the object with and without alias they should import module_and_alias_mapping = defaultdict(list) for module, aliases in self.alias_mapping.items(): module_and_alias_mapping[module].extend(aliases) for module, imports in self.module_mapping.items(): module_and_alias_mapping[module].extend( [(object, None) for object in imports] ) module_and_alias_mapping = { module: sorted(aliases) for module, aliases in module_and_alias_mapping.items() } # Now, add all of the imports we need! return updated_node.with_changes( # pyre-fixme[60]: Concatenation not yet support for multiple variadic tup... body=( *statements_before_imports, *[ parse_statement( f"from {module} import " + ", ".join( [ obj if alias is None else f"{obj} as {alias}" for (obj, alias) in aliases ] ), config=updated_node.config_for_parsing, ) for module, aliases in module_and_alias_mapping.items() if module == "__future__" ], *statements_until_add_imports, *[ parse_statement( f"import {module}", config=updated_node.config_for_parsing ) for module in sorted(self.module_imports) ], *[ parse_statement( f"import {module} as {asname}", config=updated_node.config_for_parsing, ) for (module, asname) in self.module_aliases.items() ], *[ parse_statement( f"from {module} import " + ", ".join( [ obj if alias is None else f"{obj} as {alias}" for (obj, alias) in aliases ] ), config=updated_node.config_for_parsing, ) for module, aliases in module_and_alias_mapping.items() if module != "__future__" ], *statements_after_imports, ) )
def _skip_first(orig_module: libcst.Module) -> bool: # Is there a __strict__ flag or docstring at the top? if m.matches( orig_module, m.Module( body=[ m.SimpleStatementLine( body=[ m.Assign(targets=[m.AssignTarget(target=m.Name("__strict__"))]) ] ), m.ZeroOrMore(), ] ) | m.Module( body=[ m.SimpleStatementLine(body=[m.Expr(value=m.SimpleString())]), m.ZeroOrMore(), ] ), ): return True return False