AddImportsVisitor: add imports before the first non-import statement

libcst/codemod/visitors/_add_imports.py

@@ -8,11 +8,51 @@
 
 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 GatherImportsVisitor
+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)
 
 
 class AddImportsVisitor(ContextAwareTransformer):
@@ -169,12 +209,12 @@ def __init__(
  for module in sorted(from_imports_aliases)
  }
 
- # Track the list of imports found in the file
+ # 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
- gatherer = GatherImportsVisitor(self.context)
+ # 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
 
@@ -213,6 +253,10 @@ def leave_ImportFrom(
  # 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
@@ -260,39 +304,26 @@ def _split_module(
  statement_before_import_location = 0
  import_add_location = 0
 
- # never insert an import before initial __strict__ flag
- if m.matches(
- orig_module,
- m.Module(
- body=[
- m.SimpleStatementLine(
- body=[
- m.Assign(
- targets=[m.AssignTarget(target=m.Name("__strict__"))]
- )
- ]
- ),
- m.ZeroOrMore(),
- ]
- ),
- ):
- statement_before_import_location = import_add_location = 1
-
  # 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.
- for i, statement in enumerate(orig_module.body):
- if i == 0 and m.matches(
- statement, m.SimpleStatementLine(body=[m.Expr(value=m.SimpleString())])
+
+ # 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()])
  ):
- statement_before_import_location = import_add_location = 1
- elif isinstance(statement, libcst.SimpleStatementLine):
- for possible_import in statement.body:
- for last_import in self.all_imports:
- if possible_import is last_import:
- import_add_location = i + 1
- break
+ import_add_location = i + statement_before_import_location + 1
+ else:
+ break
 
  return (
  list(updated_module.body[:statement_before_import_location]),
@@ -414,3 +445,28 @@ def leave_Module(
  *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

libcst/codemod/visitors/_gather_imports.py

@@ -12,43 +12,9 @@
 from libcst.helpers import get_absolute_module_from_package_for_import
 
 
-class GatherImportsVisitor(ContextAwareVisitor):
+class _GatherImportsMixin(ContextAwareVisitor):
  """
- Gathers all imports in a module and stores them as attributes on the instance.
- Intended to be instantiated and passed to a :class:`~libcst.Module`
- :meth:`~libcst.CSTNode.visit` method in order to gather up information about
- imports on a module. Note that this is not a substitute for scope analysis or
- qualified name support. Please see :ref:`libcst-scope-tutorial` for a more
- robust way of determining the qualified name and definition for an arbitrary
- node.
-
- After visiting a module the following attributes will be populated:
-
- module_imports
- A sequence of strings representing modules that were imported directly, such as
- in the case of ``import typing``. Each module directly imported but not aliased
- will be included here.
- object_mapping
- A mapping of strings to sequences of strings representing modules where we
- imported objects from, such as in the case of ``from typing import Optional``.
- Each from import that was not aliased will be included here, where the keys of
- the mapping are the module we are importing from, and the value is a
- sequence of objects we are importing from the module.
- module_aliases
- A mapping of strings representing modules that were imported and aliased,
- such as in the case of ``import typing as t``. Each module imported this
- way will be represented as a key in this mapping, and the value will be
- the local alias of the module.
- alias_mapping
- A mapping of strings to sequences of tuples representing modules where we
- imported objects from and aliased using ``as`` syntax, such as in the case
- of ``from typing import Optional as opt``. Each from import that was aliased
- will be included here, where the keys of the mapping are the module we are
- importing from, and the value is a tuple representing the original object
- name and the alias.
- all_imports
- A collection of all :class:`~libcst.Import` and :class:`~libcst.ImportFrom`
- statements that were encountered in the module.
+ A Mixin class for tracking visited imports.
  """
 
  def __init__(self, context: CodemodContext) -> None:
@@ -59,15 +25,10 @@ def __init__(self, context: CodemodContext) -> None:
  # Track the aliased imports in this transform
  self.module_aliases: Dict[str, str] = {}
  self.alias_mapping: Dict[str, List[Tuple[str, str]]] = {}
- # Track all of the imports found in this transform
- self.all_imports: List[Union[libcst.Import, libcst.ImportFrom]] = []
  # Track the import for every symbol introduced into the module
  self.symbol_mapping: Dict[str, ImportItem] = {}
 
- def visit_Import(self, node: libcst.Import) -> None:
- # Track this import statement for later analysis.
- self.all_imports.append(node)
-
+ def _handle_Import(self, node: libcst.Import) -> None:
  for name in node.names:
  alias = name.evaluated_alias
  imp = ImportItem(name.evaluated_name, alias=alias)
@@ -80,10 +41,7 @@ def visit_Import(self, node: libcst.Import) -> None:
  self.module_imports.add(name.evaluated_name)
  self.symbol_mapping[name.evaluated_name] = imp
 
- def visit_ImportFrom(self, node: libcst.ImportFrom) -> None:
- # Track this import statement for later analysis.
- self.all_imports.append(node)
-
+ def _handle_ImportFrom(self, node: libcst.ImportFrom) -> None:
  # Get the module we're importing as a string.
  module = get_absolute_module_from_package_for_import(
  self.context.full_package_name, node
@@ -128,3 +86,58 @@ def visit_ImportFrom(self, node: libcst.ImportFrom) -> None:
  )
  key = ia.evaluated_alias or ia.evaluated_name
  self.symbol_mapping[key] = imp
+
+
+class GatherImportsVisitor(_GatherImportsMixin):
+ """
+ Gathers all imports in a module and stores them as attributes on the instance.
+ Intended to be instantiated and passed to a :class:`~libcst.Module`
+ :meth:`~libcst.CSTNode.visit` method in order to gather up information about
+ imports on a module. Note that this is not a substitute for scope analysis or
+ qualified name support. Please see :ref:`libcst-scope-tutorial` for a more
+ robust way of determining the qualified name and definition for an arbitrary
+ node.
+
+ After visiting a module the following attributes will be populated:
+
+ module_imports
+ A sequence of strings representing modules that were imported directly, such as
+ in the case of ``import typing``. Each module directly imported but not aliased
+ will be included here.
+ object_mapping
+ A mapping of strings to sequences of strings representing modules where we
+ imported objects from, such as in the case of ``from typing import Optional``.
+ Each from import that was not aliased will be included here, where the keys of
+ the mapping are the module we are importing from, and the value is a
+ sequence of objects we are importing from the module.
+ module_aliases
+ A mapping of strings representing modules that were imported and aliased,
+ such as in the case of ``import typing as t``. Each module imported this
+ way will be represented as a key in this mapping, and the value will be
+ the local alias of the module.
+ alias_mapping
+ A mapping of strings to sequences of tuples representing modules where we
+ imported objects from and aliased using ``as`` syntax, such as in the case
+ of ``from typing import Optional as opt``. Each from import that was aliased
+ will be included here, where the keys of the mapping are the module we are
+ importing from, and the value is a tuple representing the original object
+ name and the alias.
+ all_imports
+ A collection of all :class:`~libcst.Import` and :class:`~libcst.ImportFrom`
+ statements that were encountered in the module.
+ """
+
+ 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 visit_Import(self, node: libcst.Import) -> None:
+ # Track this import statement for later analysis.
+ self.all_imports.append(node)
+ self._handle_Import(node)
+
+ def visit_ImportFrom(self, node: libcst.ImportFrom) -> None:
+ # Track this import statement for later analysis.
+ self.all_imports.append(node)
+ self._handle_ImportFrom(node)