ml_switcheroo.testing.linter

Structural Validation Linter (The Anti-Pollution Check).

This module provides the StructuralLinter, a static analysis tool designed to verify that the output of a transpilation contains no artifacts from the source framework.

Detection Scope: 1. Direct Imports: Flags import torch or from flax import …. 2. Aliased Usage: Tracks import torch as t and flags subsequent calls like t.abs(). 3. Attribute Access: Flags torch.nn.Linear if torch is a forbidden root.

Classes

StructuralLinter

Scans CST for forbidden framework usage.

Functions

validate_transpilation(→ Tuple[bool, List[str]])

Facade to lint generated code against a specific source framework.

Module Contents

class ml_switcheroo.testing.linter.StructuralLinter(forbidden_roots: Set[str])

Bases: libcst.CSTVisitor

Scans CST for forbidden framework usage.

forbidden_roots

The set of root package names (e.g., {‘torch’, ‘flax’}) that should NOT appear in the generated code.

Type:

Set[str]

violations

usage violations found during the scan.

Type:

List[str]

_local_aliases

Maps local variable names to forbidden roots. e.g., {‘t’: ‘torch’, ‘nn’: ‘torch.nn’}.

Type:

Dict[str, str]

forbidden_roots

violations: List[str] = []

check(code: str) → List[str]

Runs the linter on a source string.

Parameters:

code – The Python source code to validate.

Returns:

A list of error messages. Empty if valid.

Return type:

List[str]

visit_Import(node: libcst.Import) → None

Checks import x, import x as y.

leave_Import(node: libcst.Import) → None

visit_ImportFrom(node: libcst.ImportFrom) → None

Checks from x import y.

leave_ImportFrom(node: libcst.ImportFrom) → None

visit_Name(node: libcst.Name) → None

Checks usage of aliased forbidden variables (e.g. t.abs() where t is torch).

visit_Attribute(node: libcst.Attribute) → None

Checks attributes to provide more specific error messages (e.g. torch.abs).

ml_switcheroo.testing.linter.validate_transpilation(code: str, source_fw: str) → Tuple[bool, List[str]]

Facade to lint generated code against a specific source framework.

It automatically expands the source_fw into a set of forbidden roots including: 1. The framework itself (e.g. ‘torch’). 2. Known aliases (from adapter). 3. Inherited parents (e.g. ‘flax_nnx’ bans ‘jax’ too).

Parameters:

• code – The generated python code.

• source_fw – The framework that SHOULD have been removed (e.g., “torch”).

Returns:

(is_valid, list_of_errors)