ml_switcheroo.core.import_fixer

Import Management Transformer for ml-switcheroo.

This module is responsible for reading the AST’s import references and remapping them based on the data provided by SemanticsManager (Feature 024). It also handles collapsing fully qualified paths (e.g. torch.nn.Linear) into aliases (e.g. nn.Linear) based on content injection logic.

Classes

ImportFixer

LibCST Transformer that manages top-level imports and path shortening.

Module Contents

class ml_switcheroo.core.import_fixer.ImportFixer(source_fws: str | List[str], target_fw: str, submodule_map: Dict[str, Tuple[str, str | None, str | None]], alias_map: Dict[str, Tuple[str, str]] | None = None, preserve_source: bool = False)

Bases: libcst.CSTTransformer

LibCST Transformer that manages top-level imports and path shortening.

target_fw

submodule_map

alias_map

preserve_source = False

leave_Attribute(original_node: libcst.Attribute, updated_node: libcst.Attribute) → libcst.BaseExpression

Collapses fully qualified paths to aliases if configured. e.g. torch.nn.Linear -> nn.Linear if torch.nn -> nn alias exists for this target.

FIX: Uses original_node to determine the full path. This prevents greedy collapsing of shorter prefixes (e.g. ‘torch.nn’ -> ‘nn’) from breaking matching of longer prefixes (e.g. ‘torch.nn.functional’ -> ‘F’) when traversal bubbles up.

leave_Import(original_node: libcst.Import, updated_node: libcst.Import) → libcst.Import | libcst.RemovalSentinel

Inspects import ... statements. Checks aliases against submodule_map or prunes matches to source_fw.

leave_ImportFrom(original_node: libcst.ImportFrom, updated_node: libcst.ImportFrom) → libcst.ImportFrom | libcst.RemovalSentinel

Inspects from ... import ... statements. Checks module+name against submodule_map.

leave_Module(original_node: libcst.Module, updated_node: libcst.Module) → libcst.Module

Post-process module to inject imports intelligently.