Python: imports in an if TYPE_CHECKING: block are module scope

rewrite-python/rewrite/src/rewrite/python/add_import.py

@@ -253,8 +253,8 @@ def _try_merge_into_existing(self, cu: CompilationUnit) -> CompilationUnit:
             # case-insensitive alphabetical position (mirrors rewrite-javascript's
             # AddImport). Existing members are not reordered; only the new member
             # is positioned.
-            new_import = self._create_import_element(self.name, self.alias)
-            existing_padded = self._insert_member(
+            new_import = create_import_element(self.name, self.alias)
+            existing_padded = insert_member(
                 list(stmt.padding.names.padding.elements), new_import)
 
             # Recreate the MultiImport with the new names
@@ -279,66 +279,9 @@ def _try_merge_into_existing(self, cu: CompilationUnit) -> CompilationUnit:
 
         return cu
 
-    def _insert_member(self, elements: list, new_import: Import) -> list:
-        """Insert ``new_import`` into a list of ``JRightPadded[Import]`` at its
-        case-insensitive alphabetical position.
-
-        The space after the ``import`` keyword lives in the surrounding
-        ``JContainer.before``, so the element at index 0 carries an empty prefix
-        while every later element carries a single-space prefix (the space after
-        the separating comma). Trailing whitespace (e.g. before a ``)`` in a
-        parenthesized import) lives in the last element's ``.after`` and must
-        travel with whichever element ends up last.
-        """
-        insert_idx = self._sorted_insert_index(elements, new_import)
-        end = len(elements)
-
-        if insert_idx == 0:
-            prefix = Space.EMPTY
-            if elements:
-                # The displaced first element now follows a comma.
-                first = elements[0]
-                elements[0] = first.replace(
-                    _element=first.element.replace(prefix=Space.SINGLE_SPACE))
-        else:
-            prefix = Space.SINGLE_SPACE
-        new_import = new_import.replace(prefix=prefix)
-
-        if insert_idx == end and elements:
-            # Appending at the end: the new element becomes the last, so any
-            # trailing whitespace moves from the old last element onto it.
-            last = elements[-1]
-            elements[-1] = last.replace(_after=Space.EMPTY)
-            after = last.after
-        else:
-            after = Space.EMPTY
-
-        elements.insert(insert_idx, JRightPadded(new_import, after, Markers.EMPTY))
-        return elements
-
-    def _sorted_insert_index(self, elements: list, new_import: Import) -> int:
-        """Return the index at which the new member keeps the list in
-        case-insensitive alphabetical order: the first existing member whose
-        bound name sorts after the new member's, or the end if none does.
-
-        Members are sorted by their bound name (alias if present, else the
-        imported name), matching rewrite-javascript's comparator.
-        """
-        new_key = self._sort_key(new_import)
-        for i, padded in enumerate(elements):
-            if new_key < self._sort_key(padded.element):
-                return i
-        return len(elements)
-
-    @staticmethod
-    def _sort_key(imp: Import) -> str:
-        """Case-insensitive sort key for an imported member: its alias if it has
-        one, otherwise the imported name."""
-        return (get_alias_name(imp) or get_qualid_name(imp.qualid)).lower()
-
     def _add_import(self, cu: CompilationUnit) -> CompilationUnit:
         """Add a new import statement to the compilation unit."""
-        new_import = self._create_multi_import()
+        new_import = create_import_statement(self.module, self.name, self.alias)
 
         # Insert after the module docstring (which must stay first) and any existing imports.
         padded_stmts = list(cu.padding.statements)
@@ -383,133 +326,191 @@ def _add_import(self, cu: CompilationUnit) -> CompilationUnit:
 
         return cu.padding.replace(_statements=padded_stmts)
 
-    def _create_multi_import(self) -> MultiImport:
-        """Create a new MultiImport statement."""
-        if self.name is None:
-            # Direct import: import module [as alias]
-            import_elem = self._create_import_element(self.module, self.alias)
-            return MultiImport(
-                random_id(),
-                Space([], '\n'),
-                Markers.EMPTY,
-                None,  # No 'from'
-                False,  # Not parenthesized
-                JContainer(
-                    Space.SINGLE_SPACE,
-                    [pad_right(import_elem)],
-                    Markers.EMPTY
-                )
-            )
-        else:
-            # From import: from module import name [as alias]
-            from_name = self._create_module_name(self.module)
-            # Add space prefix (the space between 'from' and module name)
-            from_name = from_name.replace(prefix=Space.SINGLE_SPACE)
-            import_elem = self._create_import_element(self.name, self.alias)
-            return MultiImport(
-                random_id(),
-                Space([], '\n'),
-                Markers.EMPTY,
-                JRightPadded(from_name, Space.SINGLE_SPACE, Markers.EMPTY),
-                False,  # Not parenthesized
-                JContainer(
-                    Space.SINGLE_SPACE,
-                    [pad_right(import_elem)],
-                    Markers.EMPTY
-                )
-            )
 
-    def _create_import_element(self, name: str, alias: Optional[str]) -> Import:
-        """Create an Import element."""
-        qualid = self._create_qualified_name(name)
-        alias_left_padded = None
-        if alias:
-            alias_ident = Identifier(
-                random_id(),
+def insert_member(elements: list, new_import: Import) -> list:
+    """Insert ``new_import`` into a list of ``JRightPadded[Import]`` at its
+    case-insensitive alphabetical position.
+
+    The space after the ``import`` keyword lives in the surrounding
+    ``JContainer.before``, so the element at index 0 carries an empty prefix
+    while every later element carries a single-space prefix (the space after
+    the separating comma). Trailing whitespace (e.g. before a ``)`` in a
+    parenthesized import) lives in the last element's ``.after`` and must
+    travel with whichever element ends up last.
+    """
+    insert_idx = _sorted_insert_index(elements, new_import)
+    end = len(elements)
+
+    if insert_idx == 0:
+        prefix = Space.EMPTY
+        if elements:
+            # The displaced first element now follows a comma.
+            first = elements[0]
+            elements[0] = first.replace(
+                _element=first.element.replace(prefix=Space.SINGLE_SPACE))
+    else:
+        prefix = Space.SINGLE_SPACE
+    new_import = new_import.replace(prefix=prefix)
+
+    if insert_idx == end and elements:
+        # Appending at the end: the new element becomes the last, so any
+        # trailing whitespace moves from the old last element onto it.
+        last = elements[-1]
+        elements[-1] = last.replace(_after=Space.EMPTY)
+        after = last.after
+    else:
+        after = Space.EMPTY
+
+    elements.insert(insert_idx, JRightPadded(new_import, after, Markers.EMPTY))
+    return elements
+
+def _sorted_insert_index(elements: list, new_import: Import) -> int:
+    """Return the index at which the new member keeps the list in
+    case-insensitive alphabetical order: the first existing member whose
+    bound name sorts after the new member's, or the end if none does.
+
+    Members are sorted by their bound name (alias if present, else the
+    imported name), matching rewrite-javascript's comparator.
+    """
+    new_key = _sort_key(new_import)
+    for i, padded in enumerate(elements):
+        if new_key < _sort_key(padded.element):
+            return i
+    return len(elements)
+
+def _sort_key(imp: Import) -> str:
+    """Case-insensitive sort key for an imported member: its alias if it has
+    one, otherwise the imported name."""
+    return (get_alias_name(imp) or get_qualid_name(imp.qualid)).lower()
+
+
+def create_import_statement(module: str, name: Optional[str] = None,
+                            alias: Optional[str] = None) -> MultiImport:
+    """Build ``import module [as alias]`` or ``from module import name [as alias]``."""
+    if name is None:
+        # Direct import: import module [as alias]
+        import_elem = create_import_element(module, alias)
+        return MultiImport(
+            random_id(),
+            Space([], '\n'),
+            Markers.EMPTY,
+            None,  # No 'from'
+            False,  # Not parenthesized
+            JContainer(
                 Space.SINGLE_SPACE,
-                Markers.EMPTY,
-                [],
-                alias,
-                None,
-                None
+                [pad_right(import_elem)],
+                Markers.EMPTY
             )
-            alias_left_padded = JLeftPadded(
+        )
+    else:
+        # From import: from module import name [as alias]
+        from_name = _create_module_name(module)
+        # Add space prefix (the space between 'from' and module name)
+        from_name = from_name.replace(prefix=Space.SINGLE_SPACE)
+        import_elem = create_import_element(name, alias)
+        return MultiImport(
+            random_id(),
+            Space([], '\n'),
+            Markers.EMPTY,
+            JRightPadded(from_name, Space.SINGLE_SPACE, Markers.EMPTY),
+            False,  # Not parenthesized
+            JContainer(
                 Space.SINGLE_SPACE,
-                alias_ident,
+                [pad_right(import_elem)],
                 Markers.EMPTY
             )
+        )
 
-        return Import(
+def create_import_element(name: str, alias: Optional[str] = None) -> Import:
+    """Build one member of an import statement: ``name [as alias]``."""
+    qualid = _create_qualified_name(name)
+    alias_left_padded = None
+    if alias:
+        alias_ident = Identifier(
             random_id(),
-            Space.EMPTY,
+            Space.SINGLE_SPACE,
             Markers.EMPTY,
-            JLeftPadded(Space.EMPTY, False, Markers.EMPTY),
-            qualid,
-            alias_left_padded
+            [],
+            alias,
+            None,
+            None
+        )
+        alias_left_padded = JLeftPadded(
+            Space.SINGLE_SPACE,
+            alias_ident,
+            Markers.EMPTY
         )
 
-    def _create_qualified_name(self, name: str) -> FieldAccess:
-        """Create a FieldAccess for a qualified name."""
-        parts = name.split('.')
-        if len(parts) == 1:
-            return FieldAccess(
-                random_id(),
+    return Import(
+        random_id(),
+        Space.EMPTY,
+        Markers.EMPTY,
+        JLeftPadded(Space.EMPTY, False, Markers.EMPTY),
+        qualid,
+        alias_left_padded
+    )
+
+def _create_qualified_name(name: str) -> FieldAccess:
+    """Create a FieldAccess for a qualified name."""
+    parts = name.split('.')
+    if len(parts) == 1:
+        return FieldAccess(
+            random_id(),
+            Space.EMPTY,
+            Markers.EMPTY,
+            Empty(random_id(), Space.EMPTY, Markers.EMPTY),
+            JLeftPadded(
                 Space.EMPTY,
-                Markers.EMPTY,
-                Empty(random_id(), Space.EMPTY, Markers.EMPTY),
-                JLeftPadded(
-                    Space.EMPTY,
-                    Identifier(random_id(), Space.EMPTY, Markers.EMPTY, [], parts[0], None, None),
-                    Markers.EMPTY
-                ),
-                None
-            )
+                Identifier(random_id(), Space.EMPTY, Markers.EMPTY, [], parts[0], None, None),
+                Markers.EMPTY
+            ),
+            None
+        )
 
-        # Build nested FieldAccess for qualified names like "os.path"
-        # Start with the first part as an Identifier target
-        result: J = Identifier(random_id(), Space.EMPTY, Markers.EMPTY, [], parts[0], None, None)
-        # Wrap remaining parts as FieldAccess nodes
-        for part in parts[1:]:
-            result = FieldAccess(
-                random_id(),
+    # Build nested FieldAccess for qualified names like "os.path"
+    # Start with the first part as an Identifier target
+    result: J = Identifier(random_id(), Space.EMPTY, Markers.EMPTY, [], parts[0], None, None)
+    # Wrap remaining parts as FieldAccess nodes
+    for part in parts[1:]:
+        result = FieldAccess(
+            random_id(),
+            Space.EMPTY,
+            Markers.EMPTY,
+            result,
+            JLeftPadded(
                 Space.EMPTY,
-                Markers.EMPTY,
-                result,
-                JLeftPadded(
-                    Space.EMPTY,
-                    Identifier(random_id(), Space.EMPTY, Markers.EMPTY, [], part, None, None),
-                    Markers.EMPTY
-                ),
-                None
-            )
-        # For multi-part names, result is already a FieldAccess.
-        # Wrap single-part Identifier in FieldAccess(Empty, name) for consistency
-        # (but single-part is handled above, so this shouldn't happen)
-        assert isinstance(result, FieldAccess)
-        return result
-
-    def _create_module_name(self, name: str) -> J:
-        """Create a name tree for use as a 'from' module name.
-
-        Single-part names become Identifier; multi-part become FieldAccess.
-        Unlike _create_qualified_name, this does NOT wrap single parts in
-        FieldAccess(Empty, name) since the 'from' printer has no special
-        handling for Empty targets.
-        """
-        parts = name.split('.')
-        result: J = Identifier(random_id(), Space.EMPTY, Markers.EMPTY, [], parts[0], None, None)
-        for part in parts[1:]:
-            result = FieldAccess(
-                random_id(),
+                Identifier(random_id(), Space.EMPTY, Markers.EMPTY, [], part, None, None),
+                Markers.EMPTY
+            ),
+            None
+        )
+    # For multi-part names, result is already a FieldAccess.
+    # Wrap single-part Identifier in FieldAccess(Empty, name) for consistency
+    # (but single-part is handled above, so this shouldn't happen)
+    assert isinstance(result, FieldAccess)
+    return result
+
+def _create_module_name(name: str) -> J:
+    """Create a name tree for use as a 'from' module name.
+
+    Single-part names become Identifier; multi-part become FieldAccess.
+    Unlike _create_qualified_name, this does NOT wrap single parts in
+    FieldAccess(Empty, name) since the 'from' printer has no special
+    handling for Empty targets.
+    """
+    parts = name.split('.')
+    result: J = Identifier(random_id(), Space.EMPTY, Markers.EMPTY, [], parts[0], None, None)
+    for part in parts[1:]:
+        result = FieldAccess(
+            random_id(),
+            Space.EMPTY,
+            Markers.EMPTY,
+            result,
+            JLeftPadded(
                 Space.EMPTY,
-                Markers.EMPTY,
-                result,
-                JLeftPadded(
-                    Space.EMPTY,
-                    Identifier(random_id(), Space.EMPTY, Markers.EMPTY, [], part, None, None),
-                    Markers.EMPTY
-                ),
-                None
-            )
-        return result
-
+                Identifier(random_id(), Space.EMPTY, Markers.EMPTY, [], part, None, None),
+                Markers.EMPTY
+            ),
+            None
+        )
+    return result