Custom Bytecode & Metaprogramming

Overview

Advanced techniques for code transformation and generation: - AST manipulation: Transform code before compilation - Code generation: Generate code programmatically - Decorators: Intercept and modify function behavior - Metaclasses: Control how classes are created - Dynamic methods: Create functions at runtime - DSLs: Build domain-specific languages

Understanding these enables building powerful ML frameworks.

---

AST Transformation Patterns

Pattern 1: Add Tracing to Functions

import ast
import inspect
from functools import wraps

class TracingTransformer(ast.NodeTransformer):
    """Add print statements to trace execution."""

    def visit_FunctionDef(self, node):
        # Add trace at function entry
        trace_enter = ast.Expr(
            value=ast.Call(
                func=ast.Name(id='print', ctx=ast.Load()),
                args=[ast.Constant(value=f'Entering {node.name}')],
                keywords=[]
            )
        )

        # Add trace before return statements
        new_body = [trace_enter]
        for stmt in node.body:
            if isinstance(stmt, ast.Return):
                trace_exit = ast.Expr(
                    value=ast.Call(
                        func=ast.Name(id='print', ctx=ast.Load()),
                        args=[ast.Constant(value=f'Exiting {node.name}')],
                        keywords=[]
                    )
                )
                new_body.append(trace_exit)
            new_body.append(stmt)

        node.body = new_body
        return node

# Example function
code = """
def add(a, b):
    result = a + b
    return result
"""

# Transform and execute
tree = ast.parse(code)
transformer = TracingTransformer()
new_tree = transformer.visit(tree)
compiled = compile(new_tree, '<ast>', 'exec')
exec(compiled)

add(3, 5)
# Output:
# Entering add
# Exiting add

Pattern 2: Auto-Vectorization with AST

import ast
import numpy as np

class VectorizationTransformer(ast.NodeTransformer):
    """Convert element-wise loops to NumPy operations."""

    def visit_For(self, node):
        # Detect pattern: for i in range(n):
        #                     result[i] = operation(array[i])

        # For simplicity, just visit children
        self.generic_visit(node)
        return node

# Example: Simple addition loop
def add_arrays_loop(a, b):
    """Add two arrays with loop."""
    result = []
    for i in range(len(a)):
        result.append(a[i] + b[i])
    return result

# Vectorized version
def add_arrays_vectorized(a, b):
    """Add two arrays with NumPy."""
    return a + b

# Results are identical
a = np.array([1, 2, 3])
b = np.array([4, 5, 6])
print(add_arrays_loop(a.tolist(), b.tolist()))    # [5, 7, 9]
print(add_arrays_vectorized(a, b))                # [5 7 9]

---

Code Generation Patterns

Pattern 1: Generate Specialized Functions

def generate_adder(n):
    """Generate a function that adds n to its input."""
    code = f"""
def adder(x):
    return x + {n}
"""
    namespace = {}
    exec(code, namespace)
    return namespace['adder']

# Create specialized functions
add_5 = generate_adder(5)
add_10 = generate_adder(10)

print(add_5(3))   # 8
print(add_10(3))  # 13

Pattern 2: Dynamic Method Creation

class DynamicModel:
    """Create methods dynamically based on configuration."""

    def __init__(self, config):
        self.config = config
        self._build_layers()

    def _build_layers(self):
        """Dynamically create layer methods."""
        for i, layer_size in enumerate(self.config):
            # Create weight matrix
            weights = np.random.randn(layer_size, layer_size)

            # Create forward function
            def forward(x, w=weights):
                return np.dot(x, w)

            # Add to class
            setattr(self, f'layer_{i}', forward)

    def forward(self, x):
        """Apply all layers."""
        for i in range(len(self.config)):
            layer = getattr(self, f'layer_{i}')
            x = layer(x)
        return x

# Use it
import numpy as np
config = [100, 50, 10]
model = DynamicModel(config)
output = model.forward(np.random.randn(100))
print(output.shape)  # (10,)

---

Decorator-Based Code Transformation

Pattern 1: Caching Decorator

from functools import wraps

def memoize(func):
    """Cache function results."""
    cache = {}

    @wraps(func)
    def wrapper(*args, **kwargs):
        key = (args, tuple(sorted(kwargs.items())))

        if key not in cache:
            cache[key] = func(*args, **kwargs)

        return cache[key]

    return wrapper

@memoize
def fibonacci(n):
    if n <= 1:
        return n
    return fibonacci(n-1) + fibonacci(n-2)

# Much faster due to caching
print(fibonacci(35))  # Instant (with cache)

Pattern 2: Type Checking Decorator

from functools import wraps
import inspect

def typed(func):
    """Enforce type hints at runtime."""
    sig = inspect.signature(func)

    @wraps(func)
    def wrapper(*args, **kwargs):
        # Bind arguments to parameters
        bound = sig.bind(*args, **kwargs)
        bound.apply_defaults()

        # Check types
        for param_name, param_value in bound.arguments.items():
            param = sig.parameters[param_name]

            if param.annotation != inspect.Parameter.empty:
                expected_type = param.annotation
                if not isinstance(param_value, expected_type):
                    raise TypeError(
                        f"{param_name} must be {expected_type.__name__}, "
                        f"got {type(param_value).__name__}"
                    )

        return func(*args, **kwargs)

    return wrapper

@typed
def add_numbers(a: int, b: int) -> int:
    return a + b

print(add_numbers(3, 5))       # 8
# print(add_numbers(3, "5"))   # TypeError!

---

Metaclasses

Pattern 1: Automatic Method Registration

class RegistryMeta(type):
    """Metaclass that auto-registers methods."""

    def __new__(cls, name, bases, namespace):
        namespace['_registry'] = {}

        for key, value in namespace.items():
            if callable(value) and key.startswith('handle_'):
                # Register handler
                event_name = key.replace('handle_', '')
                namespace['_registry'][event_name] = value

        return super().__new__(cls, name, bases, namespace)

class EventHandler(metaclass=RegistryMeta):
    """Automatically registers event handlers."""

    def handle_click(self, event):
        print(f"Click: {event}")

    def handle_hover(self, event):
        print(f"Hover: {event}")

# Check registry
print(EventHandler._registry)
# {'click': <function>, 'hover': <function>}

Pattern 2: Declarative API with Metaclass

class Field:
    def __init__(self, field_type):
        self.field_type = field_type

class ModelMeta(type):
    """Metaclass for declarative model definitions."""

    def __new__(cls, name, bases, namespace):
        # Collect fields
        fields = {}
        for key, value in namespace.items():
            if isinstance(value, Field):
                fields[key] = value

        namespace['_fields'] = fields
        return super().__new__(cls, name, bases, namespace)

class Model(metaclass=ModelMeta):
    """Base for declarative models."""
    pass

class User(Model):
    """Declarative model definition."""
    name = Field(str)
    age = Field(int)
    email = Field(str)

# Access fields
print(User._fields)
# {'name': Field(str), 'age': Field(int), 'email': Field(str)}

---

Dynamic Import Hooks

Pattern: Custom Module Loading

import sys
import importlib.abc
import importlib.machinery

class CustomLoader(importlib.abc.Loader):
    """Custom loader that transforms code."""

    def exec_module(self, module):
        """Execute module code with transformation."""
        source = module.__loader__.get_source(module.__name__)

        # Transform source
        tree = ast.parse(source)
        # ... apply transformations ...

        # Compile and execute
        code = compile(tree, module.__file__, 'exec')
        exec(code, module.__dict__)

# Register custom loader
# sys.meta_path.insert(0, CustomFinder())

---

Real-World: PyTorch's autograd DSL

How PyTorch Uses AST/Metaprogramming

# Simplified version of how PyTorch works

class Tensor:
    """Simplified Tensor with gradient tracking."""

    def __init__(self, data, requires_grad=False):
        self.data = data
        self.requires_grad = requires_grad
        self.grad = None
        self._grad_fn = None  # Computation graph node

    def __add__(self, other):
        """Intercept addition operator."""
        result = Tensor(self.data + other.data, requires_grad=True)
        result._grad_fn = ('add', self, other)
        return result

    def __mul__(self, other):
        """Intercept multiplication operator."""
        result = Tensor(self.data * other.data, requires_grad=True)
        result._grad_fn = ('mul', self, other)
        return result

    def backward(self, grad=1.0):
        """Backpropagation through computation graph."""
        if self._grad_fn is None:
            return

        op, a, b = self._grad_fn

        if op == 'add':
            a.grad = grad
            b.grad = grad
        elif op == 'mul':
            a.grad = grad * b.data
            b.grad = grad * a.data

        if a.requires_grad:
            a.backward(a.grad)
        if b.requires_grad:
            b.backward(b.grad)

# Use it
x = Tensor([1, 2, 3], requires_grad=True)
y = Tensor([4, 5, 6], requires_grad=True)

z = x * y + x  # Builds computation graph
z.backward()

print(f"x.grad: {x.grad}")  # [4, 5, 6] (from y) + 1 (from addition)
print(f"y.grad: {y.grad}")  # [1, 2, 3] (from x)

---

Summary: Metaprogramming Techniques

Technique	Use Case	Complexity
AST Transformation	Code optimization, tracing	High
Code Generation	Specialized functions, DSLs	High
Decorators	Caching, type checking, timing	Low
Metaclasses	Auto-registration, declarative APIs	High
Import Hooks	Custom module loading	Very High
Dynamic Methods	Runtime API creation	Medium

---

Best Practices

1. Keep it simple: Only use metaprogramming when necessary
2. Document heavily: Advanced techniques are hard to understand
3. Profile first: Measure before optimizing with metaprogramming
4. Test thoroughly: Dynamically generated code is harder to test
5. Use type hints: Help IDE and static analysis understand code

---

Related Topics

• 01 Python Bytecode Fundamentals - Understanding generated bytecode
• 02 Execution Model & Compilation - How transformed code is compiled
• 03 Jit Compilation & Optimization - Optimizing generated code
• 02 Decorators - Decorator patterns