ml_switcheroo.core.rewriter.contexts

ContextRewriter Mixin.

This module provides the ContextMixin for the PivotRewriter, handling transformations of With blocks (Context Managers).

Capabilities: 1. API Mapping: Maps context managers (e.g., torch.no_grad() -> contextlib.nullcontext()). 2. Structural Stripping: Supports completely removing the with block while preserving the

body indentation (Lifting), if the semantic mapping dictates transformation_type=”strip_context”.

Classes

ContextMixin

Mixin for transforming With blocks.

Module Contents

class ml_switcheroo.core.rewriter.contexts.ContextMixin(semantics: ml_switcheroo.semantics.manager.SemanticsManager, config: ml_switcheroo.config.RuntimeConfig)

Bases: ml_switcheroo.core.rewriter.base.BaseRewriter

Mixin for transforming With blocks.

leave_With(original_node: libcst.With, updated_node: libcst.With) → libcst.With | libcst.FlattenSentinel

Processes ‘with’ statements.

Logic: 1. Iterate over with items (expressions). 2. Identify if the expression corresponds to a CONTEXT OpType in semantics. 3. Check transformation logic:

• If strip_context, identifying marker is found, lift the body.

• Otherwise, allow CallMixin (which runs on children) to have already renamed the API.