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
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'

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'

---

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](/05-py3/09-bytecode-and-execution/(02-execution-model-compilation/) - How transformed code is compiled
• [03 Jit Compilation & Optimization](/05-py3/09-bytecode-and-execution/(03-jit-compilation-optimization/) - Optimizing generated code
• 02 Decorators - Decorator patterns