ml_switcheroo.core.hooks¶

Plugin Binding Infrastructure.

This module provides the infrastructure for extending ml-switcheroo via plugins. It enables developers to intercept and modify the Abstract Syntax Tree (AST) during the conversion process using a hook-based system.

Refactor:

Added auto_wire support to the register_hook decorator.
Plugins can now declare their own Semantic definitions.
HookContext now exposes plugin_traits and current_variant for data-driven logic.
New: resolve_type method to query Symbol Table.

Attributes¶

`SymbolTableType`
`SemanticsManagerType`
`T`
`ArgInjectorType`
`PreambleInjectorType`
`HookFunction`

Classes¶

`AutoWireSpec`	Schema for plugin self-registration metadata.
`HookContext`	Context object passed to every plugin hook during transcoding.

Functions¶

`register_hook`(→ Callable[[HookFunction], HookFunction])	Decorator to register a function as a plugin hook.
`get_hook`(→ Optional[HookFunction])	Retrieves a registered hook function by its trigger name.
`get_all_hook_metadata`(→ Dict[str, AutoWireSpec])	Returns the metadata for all registered hooks.
`clear_hooks`(→ None)	Resets the internal hook registry. Primarily for testing.
`load_plugins`(→ int)	Dynamically imports plugins.

Module Contents¶

ml_switcheroo.core.hooks.SymbolTableType¶

ml_switcheroo.core.hooks.SemanticsManagerType¶

ml_switcheroo.core.hooks.T¶

ml_switcheroo.core.hooks.ArgInjectorType¶

ml_switcheroo.core.hooks.PreambleInjectorType¶

class ml_switcheroo.core.hooks.AutoWireSpec(/, **data: Any)[source]¶

Bases: pydantic.BaseModel

Schema for plugin self-registration metadata. Allows a plugin to define the Semantic Operation it satisfies.

model_config¶: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

ops: Dict[str, Dict[str, Any]] = None¶

class ml_switcheroo.core.hooks.HookContext(semantics: SemanticsManagerType, config: ml_switcheroo.config.RuntimeConfig, arg_injector: ArgInjectorType | None = None, preamble_injector: PreambleInjectorType | None = None, symbol_table: SymbolTableType | None = None)[source]¶

Context object passed to every plugin hook during transcoding.

Provides read-only access to global state and write access to specific injection points (signature args, function body preambles). Now exposes plugin_traits and current_variant for data-driven decisions.

semantics¶

source_fw¶

target_fw¶

metadata: Dict[str, Any]¶

current_op_id: str | None = None¶

resolve_type(node: Any) → str | None[source]¶

Queries the Symbol Table for the inferred type of a node.

Parameters:: node – The LibCST node to inspect.
Returns:: “Tensor” if it’s a tensor, “Module” if module, or None.
Return type:: str

property plugin_traits: ml_switcheroo.semantics.schema.PluginTraits¶

Returns the capabilities of the current Target Framework. This allows plugins to check functionality (e.g. has_numpy_compatible_arrays) rather than checking the framework name string.

Returns:: The capability flags for the target framework.
Return type:: PluginTraits

property current_variant: ml_switcheroo.semantics.schema.FrameworkVariant | None¶

Returns the Variant definition for the current operation/target. Allows plugins to read extra metadata defined in the JSON (e.g. pack_to_tuple).

Returns:: The variant definition if resolved, else None.
Return type:: Optional[FrameworkVariant]

inject_signature_arg(name: str, annotation: str | None = None) → None[source]¶

Requests injection of argument into the current function signature.

Parameters:

name (str) – The name of the argument to inject.
annotation (Optional[str]) – Type hint string for the argument.

inject_preamble(code_str: str) → None[source]¶

Requests injection of a statement at the beginning of the function body.

Parameters:: code_str (str) – Python source code string to inject.

raw_config(key: str, default: Any = None) → Any[source]¶

Retrieve a raw value from the unstructured plugin settings dict.

Parameters:

key (str) – Configuration key.
default (Any) – Default value if key is not found.

Returns:

The configuration value.

Return type:

Any

validate_settings(model: Type[T]) → T[source]¶

Validates global config against a Plugin-specific Pydantic schema.

Parameters:: model (Type[T]) – Pydantic model definition.
Returns:: Validated configuration object.
Return type:: T

lookup_api(op_name: str) → str | None[source]¶

Resolves target framework’s API string for a given standard operation.

Parameters:: op_name (str) – Standard operation ID.
Returns:: The target API string, or None if not found.
Return type:: Optional[str]

lookup_signature(op_name: str) → List[str][source]¶

Retrieves standard argument list for a given operation.

Parameters:: op_name (str) – Standard operation ID.
Returns:: List of argument names.
Return type:: List[str]

ml_switcheroo.core.hooks.HookFunction¶

ml_switcheroo.core.hooks.register_hook(trigger: str, auto_wire: Dict[str, Any] | None = None) → Callable[[HookFunction], HookFunction][source]¶

Decorator to register a function as a plugin hook.

Parameters:

trigger – The unique identifier. Can be an operation ID or a reserved system event like “transform_for_loop”.
auto_wire – Optional dictionary defining the Semantic Spec for this plugin. If provided, the SemanticsManager will automatically load these definitions, eliminating the need for JSON usage. Format matches semantics/*.json schema (e.g. {“ops”: {“MyOp”: …}}).

Returns:

The decorator wrapper.

Return type:

Callable

ml_switcheroo.core.hooks.get_hook(trigger: str) → HookFunction | None[source]¶

Retrieves a registered hook function by its trigger name. Lazily loads plugins from the default directory if registry is empty.

Parameters:: trigger (str) – Hook identifier key.
Returns:: The registered function or None.
Return type:: Optional[HookFunction]

ml_switcheroo.core.hooks.get_all_hook_metadata() → Dict[str, AutoWireSpec][source]¶

Returns the metadata for all registered hooks. Used by SemanticsManager to hydrate the Knowledge Base.

Returns:: Metadata for autowired plugins.
Return type:: Dict[str, AutoWireSpec]

ml_switcheroo.core.hooks.clear_hooks() → None[source]¶: Resets the internal hook registry. Primarily for testing.

ml_switcheroo.core.hooks.load_plugins(plugins_dir: pathlib.Path | None = None, extra_dirs: List[pathlib.Path] | None = None) → int[source]¶

Dynamically imports plugins.

Parameters:

plugins_dir – Overrides default package directory. If provided, this directory is scanned for .py files. If NOT provided, the internal ml_switcheroo.plugins package is loaded.
extra_dirs – Additional directories to scan (e.g. user extensions).

Returns:

Number of modules loaded.

Return type:

int