ml_switcheroo.core.import_fixer.attributes_mixin

Attribute Logic Mixin.

Handles the transformation of Attribute nodes, specifically for collapsing fully qualified names into aliases (e.g. jax.numpy.abs -> jnp.abs) and simplifying re-export chains.

Classes

AttributeMixin

Mixin for processing Attribute nodes.

Module Contents

class ml_switcheroo.core.import_fixer.attributes_mixin.AttributeMixin

Bases: libcst.CSTTransformer

Mixin for processing Attribute nodes.

leave_Attribute(original_node: libcst.Attribute, updated_node: libcst.Attribute) → libcst.BaseExpression

Handles path simplification and alias collapsing.

Parameters:

• original_node – The node before transformation.

• updated_node – The node after transformation.

Returns:

The transformed expression node.