ml_switcheroo.core.scanners¶

AST Scanners for Symbol Usage Detection.

This module provides LibCST visitors that analyze code to determine if specific names, frameworks, or aliases are actively referenced in the source body.

These scanners are critical for the ImportFixer logic:

If a framework alias (e.g., jnp) is injected, SimpleNameScanner verifies it is actually used before committing the import.
If a source import (e.g., import torch) is slated for removal, UsageScanner checks if it persists in the code (e.g., inside an Escape Hatch) to prevent breaking valid code.

Classes¶

`SimpleNameScanner`	Scans for the usage of a specific identifier in the code body.
`UsageScanner`	Scans the AST for usages of a specific framework root or its local aliases.

Functions¶

get_full_name(→ str)

Recursively resolves a CST Name or Attribute chain to a dot-separated string.

Module Contents¶

ml_switcheroo.core.scanners.get_full_name(node: libcst.Name | libcst.Attribute) → str¶

Recursively resolves a CST Name or Attribute chain to a dot-separated string.

This helper flattens the AST representation of dotted names into strings comparable with import definitions.

Parameters:: node – The CST node representing the identifier. Typically a cst.Name (e.g., x) or cst.Attribute (e.g., x.y).
Returns:: The fully qualified string representation (e.g., “torch.nn.functional”). Returns an empty string if the node structure is not a supported Name/Attribute chain.
Return type:: str

class ml_switcheroo.core.scanners.SimpleNameScanner(target_name: str)¶

Bases: libcst.CSTVisitor

Scans for the usage of a specific identifier in the code body.

This visitor is designed to check for the presence of variables or aliases (like jnp, tf, mx) outside of import statements. It is used to determine if a speculative import injection is actually required.

target_name¶

found = False¶

visit_Import(node: libcst.Import) → None¶: Flags entry into an import ... statement. Names appearing here are definitions, not usages.

leave_Import(node: libcst.Import) → None¶: Flags exit from an import ... statement.

visit_ImportFrom(node: libcst.ImportFrom) → None¶: Flags entry into a from ... import ... statement.

leave_ImportFrom(node: libcst.ImportFrom) → None¶: Flags exit from a from ... import ... statement.

visit_Name(node: libcst.Name) → None¶

Checks if the visited name matches the target.

If the name matches target_name and we are NOT currently inside an import definition, we mark found = True.

Parameters:: node – The name node being visited.

should_traverse(_node: libcst.CSTNode) → bool¶

Optimization hook to stop traversal once found.

Returns:: False if the target has already been found, effectively short-circuiting the rest of the AST traversal.
Return type:: bool

class ml_switcheroo.core.scanners.UsageScanner(source_fw: str)¶

Bases: libcst.CSTVisitor

Scans the AST for usages of a specific framework root or its local aliases.

This class implements a multi-pass logic during a single traversal:

Cataloging: It identifies all aliases bound to the source_fw (e.g., import torch as t -> t is an alias).
Detection: It checks if any of these cataloged aliases are used in the code body (e.g., t.abs(x)).

This is primarily used by the ImportFixer to decide whether to prune the original import statement. If usages persist (e.g., via the Escape Hatch), the import MUST be preserved to keep the code valid.

source_fw¶

found_usages: Set[str]¶

get_result() → bool¶

Returns the scan result.

Returns:: True if any tracked alias was found used in the body.
Return type:: bool

visit_Import(node: libcst.Import) → None¶

Catalogs names bound by import ....

Logic:

import torch -> tracks ‘torch’.
import torch as t -> tracks ‘t’.
import torch.nn as nn -> tracks ‘nn’ (because it stems from torch).

Parameters:: node – The import node.

leave_Import(node: libcst.Import) → None¶: Exit import scope.

visit_ImportFrom(node: libcst.ImportFrom) → None¶

Catalogs names bound by from ... import ....

Logic:

from torch import nn -> tracks ‘nn’.
from torch.nn import Linear -> tracks ‘Linear’.

Parameters:: node – The import-from node.

leave_ImportFrom(node: libcst.ImportFrom) → None¶: Exit import-from scope.

visit_Name(node: libcst.Name) → None¶

Checks if a name in the body matches one of our tracked aliases.

If found, it is recorded in found_usages.

Parameters:: node – The name node.