injector_fw

LibCST Transformer for Injecting Framework Mappings.

This module provides the logic to modify framework adapter files (e.g. torch.py) by locating the specific class registered for a framework and injecting a new StandardMap definition into its definitions property.

It handles: 1. Definitions Injection: Appending the mapping to the definitions dictionary. 2. Smart Import Injection: Analyzing the target API path (e.g. scipy.special.erf)

and injecting necessary top-level imports (import scipy) if missing.

3. Variant Parameter Injection: Supporting inject_args for adding fixed arguments.

4. Complex Literal Support: Recursively converting Lists, Tuples, and Dicts to CST nodes.

Classes

FrameworkInjector

Injects a StandardMap entry into a Framework Adapter's definitions.

Module Contents

class injector_fw.FrameworkInjector(target_fw: str, op_name: str, variant: ml_switcheroo.core.dsl.FrameworkVariant)

Bases: libcst.CSTTransformer

Injects a StandardMap entry into a Framework Adapter’s definitions.

It performs a targeted search for: 1. A class decorated with @register_framework(“target_fw”). 2. A method named definitions decorated with @property. 3. A return statement returning a Dict.

Additionally, it scans the module for existing imports and injects missing dependencies required by the new mapping (e.g. injecting import scipy if the API is scipy.special.erf).

target_fw

The framework key to look for (e.g. ‘torch’).

Type:

str

op_name

The abstract operation name (e.g. ‘LogSoftmax’).

Type:

str

variant

The configuration to inject.

Type:

FrameworkVariant

found

True if intrusion was successful.

Type:

bool

target_fw

op_name

variant

found = False

visit_Import(node: libcst.Import) → None

Tracks existing top-level imports.

visit_ImportFrom(node: libcst.ImportFrom) → None

Tracks existing from-imports.

visit_ClassDef(node: libcst.ClassDef) → bool | None

Checks if the class matches the target framework via decorator.

leave_ClassDef(original_node: libcst.ClassDef, updated_node: libcst.ClassDef) → libcst.ClassDef

Exit class scope.

visit_FunctionDef(node: libcst.FunctionDef) → bool | None

Enters functions. Only interested if inside target class and named ‘definitions’.

leave_FunctionDef(original_node: libcst.FunctionDef, updated_node: libcst.FunctionDef) → libcst.FunctionDef

Exit function scope.

leave_Return(original_node: libcst.Return, updated_node: libcst.Return) → libcst.Return | libcst.RemovalSentinel

Intercepts the return statement of the definitions property. Injects the new key-value into the dictionary.

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

Post-process module to inject missing top-level imports.