Refactors, typos and typing changes

libcst/codemod/visitors/_add_imports.py

@@ -11,19 +11,24 @@
 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(GatherImportsVisitor):
+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:]:
@@ -32,8 +37,13 @@ def leave_Module(self, original_node: libcst.Module) -> None:
                 m.SimpleStatementLine(body=[m.ImportFrom() | m.Import()]),
             ):
                 stmt = ensure_type(stmt, SimpleStatementLine)
-                imp = ensure_type(stmt.body[0], Union[ImportFrom, Import])
-                self.all_imports.append(imp)
+                # 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:
@@ -44,12 +54,6 @@ def leave_Module(self, original_node: libcst.Module) -> None:
                 imp = ensure_type(imp, ImportFrom)
                 self._handle_ImportFrom(imp)
 
-    def visit_ImportFrom(self, node: libcst.ImportFrom) -> None:
-        pass
-
-    def visit_Import(self, node: libcst.Import) -> None:
-        pass
-
 
 class AddImportsVisitor(ContextAwareTransformer):
     """
@@ -206,13 +210,13 @@ def __init__(
         }
 
         # Track the list of imports found at the top of the file
-        self.all_imports: Set[Union[libcst.Import, libcst.ImportFrom]] = set()
+        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 = set(gatherer.all_imports)
+        self.all_imports = gatherer.all_imports
 
         self.module_imports = self.module_imports - gatherer.module_imports
         for module, alias in gatherer.module_aliases.items():
@@ -249,7 +253,7 @@ def leave_ImportFrom(
             # There's nothing to do here!
             return updated_node
 
-        # Ensure this is on of the imports at the top
+        # Ensure this is one of the imports at the top
         if original_node not in self.all_imports:
             return updated_node
 

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,16 +25,9 @@ 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)
-        self._handle_Import(node)
-
     def _handle_Import(self, node: libcst.Import) -> None:
         for name in node.names:
             alias = name.evaluated_alias
@@ -82,11 +41,6 @@ def _handle_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)
-        self._handle_ImportFrom(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(
@@ -132,3 +86,58 @@ def _handle_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)