Hooks System

Summary

This chapter covers the hooks system that allows runtime customization of flatten and build operations without modifying specs. You will learn about the HookRegistry class and the four hook types: condition hooks for conditional inclusion of nodes, field transform hooks for modifying field values during extraction, pre-flatten hooks that run before a flatten operation, and post-build hooks that run after tree reconstruction.

Concepts Covered

• HookRegistry Class
• Condition Hook
• Field Transform Hook
• Pre Flatten Hook
• Post Build Hook

Prerequisites

• Chapter 1: Foundation Concepts (Decorator Pattern)
• Chapter 3: Flatten Engine (TreeFlattener Class)
• Chapter 6: Build Engine (TreeBuilder Class)

Learning Objectives

After completing this chapter, you will be able to:

• Create and configure a HookRegistry
• Write condition hooks to control which nodes are included
• Implement field transform hooks to modify values during extraction
• Register pre-flatten hooks to prepare data before flattening
• Register post-build hooks to post-process reconstructed trees
• Combine multiple hook types into a custom processing pipeline

---

Why a Hooks System?

In previous chapters you learned how TreeSpec declarations drive the flatten and build engines. A spec file describes what gets extracted and how the tree maps to tables, but it says nothing about when or whether certain parts of the tree should be included. Real-world pipelines need runtime decisions: a credit-detail section might only be relevant for premium-tier customers, a timestamp field might need reformatting before it lands in a table, or an entire tree might need pre-processing before flattening begins.

Hard-coding those decisions into the spec would defeat the purpose of a declarative configuration. Instead, TreeSpec provides a hooks system -- a set of named callback functions that the flatten and build engines invoke at well-defined points during processing. Hooks let you inject arbitrary Python logic without touching the spec itself, keeping the spec portable across environments while the hooks adapt behaviour to each deployment context.

The hooks system is built on two foundational ideas. First, it uses the Decorator Pattern, which you met in Chapter 1, to register functions under symbolic names. Second, it passes an immutable context object to every hook so that decisions can depend on runtime state such as configuration settings, the current node in the tree, or the source file being processed.

Diagram: Hook Injection Points in the Flatten-Build Pipeline

Hook Injection Points in the Flatten-Build Pipeline

Type: flowchart | sim-id: hook-injection-points
| Library: vis-network
| Status: Specified

Learning Objective: Understand where each hook type executes within the flatten-build pipeline. Bloom Level: Understand Interactions: Hover over each stage node to see which hook type fires there. Pre-flatten hooks fire before the tree walk begins; condition hooks and field transform hooks fire during the walk; post-build hooks fire after reconstruction completes.

The TreeContext Object

Before examining the HookRegistry itself, you need to understand the object that every hook receives as its primary input. A TreeContext is a frozen (immutable) dataclass that carries runtime state through the flatten or build operation. Because it is frozen, hooks cannot accidentally mutate shared state -- a critical safety property when hooks run in parallel workers.

TreeContext exposes the following fields:

• settings -- a dictionary of arbitrary key-value pairs supplied when the flattener or builder is constructed. Typical entries include batch tier identifiers, environment names, or feature flags.
• version -- an optional string identifying the schema version being processed.
• parent_chain -- a list of node names representing the path from the root to the current position in the tree.
• current_node -- the name of the node currently being processed.
• source_file -- an optional path to the file from which the current tree was loaded.
• current_row -- for build-side hooks, the dictionary representing the current table row.
• current_object -- for flatten-side hooks, the Python object currently being walked.

TreeContext provides three builder methods -- with_row(), with_node(), and with_object() -- that return a new TreeContext with one field updated while preserving all others. This functional style ensures that the context flowing through nested tree levels never leaks state back up to a parent level.

from mountainash_treespec.core.hooks.context import TreeContext

# Create a base context with runtime settings
ctx = TreeContext(settings={"BATCH_TIER": "C", "ENV": "production"})

# Derive a child context as you descend into a node
ctx_child = ctx.with_node("CreditDetail")
print(ctx_child.current_node)   # "CreditDetail"
print(ctx_child.parent_chain)   # ["CreditDetail"]
print(ctx.current_node)         # None -- original is unchanged

The code above creates a base context carrying two settings entries and then derives a child context for a node named "CreditDetail". Because TreeContext is frozen, the original ctx remains unmodified, which prevents accidental cross-contamination between branches of the tree walk.

Field	Type	Default	Typical Use
settings	dict[str, Any]	{}	Feature flags, tier identifiers
version	Optional[str]	None	Schema version routing
parent_chain	list[str]	[]	Ancestry-based filtering
current_node	Optional[str]	None	Node-specific logic
source_file	Optional[str]	None	Per-file transforms
current_row	Optional[dict]	None	Row data during build
current_object	Optional[Any]	None	Object data during flatten

The HookRegistry Class

A HookRegistry is the central container for all hooks that a flatten or build operation can invoke. It maintains four internal dictionaries, one for each hook type, mapping symbolic names to callable functions. When you construct a TreeFlattener or TreeBuilder, you pass in a HookRegistry instance; if you omit it, the engine creates an empty registry so that specs referencing no hooks work without any additional configuration.

The registry exposes four decorator methods for registration and two dispatch methods for invocation. The decorator approach means that defining a hook and registering it are a single step -- there is no separate "add to registry" call that could be forgotten.

from mountainash_treespec.core.hooks.registry import HookRegistry

hooks = HookRegistry()

<!-- concept:97 -->
# Register a condition hook named "tier_c"
@hooks.condition("tier_c")
def tier_c(ctx):
    return ctx.settings.get("BATCH_TIER") == "C"

# Check registration
print(hooks.has_hook("tier_c"))    # True
print(hooks.hook_names())          # {"tier_c"}

The snippet above creates a registry, registers a single condition hook, and then queries the registry to confirm registration. The has_hook() method searches across all four hook dictionaries, while hook_names() returns the union of all registered names regardless of type.

Diagram: HookRegistry Internal Architecture

HookRegistry Internal Architecture

Type: class diagram | sim-id: hook-registry-architecture
| Library: vis-network
| Status: Specified

Learning Objective: Visualize the four internal dictionaries and the decorator-based registration pattern. Bloom Level: Understand Interactions: Click on each dictionary node (conditions, transforms, pre_flatten, post_build) to see its type signature and example entries. Click on decorator methods to see the registration flow.

Parallel Safety and Picklability

When you use batch processing with ProcessPoolExecutor, Python serializes objects using the pickle protocol to send them to worker processes. Lambda functions and closures that capture local variables are not picklable, which means they cannot cross process boundaries. The HookRegistry provides a validate_parallel() method that attempts to pickle every registered hook and raises an UnpicklableHookError if any hook fails serialization.

The following code demonstrates the validation mechanism. A hook defined as a module-level named function passes validation, while a lambda does not.

from mountainash_treespec.core.hooks.registry import (
    HookRegistry, UnpicklableHookError
)

hooks = HookRegistry()

# Module-level named functions ARE picklable
@hooks.condition("always_true")
def always_true(ctx):
    return True

hooks.validate_parallel()  # passes silently

# Lambdas are NOT picklable
hooks._conditions["bad_hook"] = lambda ctx: True
try:
    hooks.validate_parallel()
except UnpicklableHookError as e:
    print(f"Caught: {e}")
    # "Hook 'bad_hook' is not picklable (required for parallel mode): ..."

As a rule, always define hooks as named functions at module scope when parallel processing is a possibility. This constraint also encourages separating hook logic into dedicated modules, which improves testability.

Condition Hooks

A condition hook is a function that receives a TreeContext and returns a boolean. During the tree walk, whenever the flattener encounters a ChildSpec that has a condition attribute, it calls evaluate_condition() on the registry with the given name and the current context. If the function returns True, the child is processed normally. If it returns False, the child and all its descendants are skipped entirely -- no columns appear in the output table for that subtree.

This mechanism is powerful because it lets a single spec definition cover multiple scenarios. A financial data spec might include both standard and premium fields; a condition hook decides at runtime which tier applies, and the same spec file works for both without modification.

Consider a concrete example. An Account dataclass has an optional CreditDetail child, and the spec marks that child with condition="tier_c":

from dataclasses import dataclass
from typing import Optional

@dataclass
class CreditDetail:
    score: Optional[int] = None
    provider: Optional[str] = None

@dataclass
class Account:
    account_id: str = ""
    holder_name: str = ""
    credit_detail: Optional[CreditDetail] = None

The spec configuration marks the credit_detail child with a condition name. Then a HookRegistry provides the logic that decides when to include it.

from mountainash_treespec.core.hooks.registry import HookRegistry
from mountainash_treespec.core.hooks.context import TreeContext
from mountainash_treespec.core.flatten.flattener import TreeFlattener

hooks = HookRegistry()

@hooks.condition("tier_c")
def tier_c(ctx: TreeContext) -> bool:
    return ctx.settings.get("BATCH_TIER") == "C"

# Tier C: credit detail IS included
flattener_c = TreeFlattener(spec, hooks=hooks, settings={"BATCH_TIER": "C"})
tables_c = flattener_c.flatten(account)
assert "Credit_score" in tables_c["accounts"][0]

# Tier N: credit detail is EXCLUDED
flattener_n = TreeFlattener(spec, hooks=hooks, settings={"BATCH_TIER": "N"})
tables_n = flattener_n.flatten(account)
assert "Credit_score" not in tables_n["accounts"][0]

The same spec, the same data object, and the same hook registry produce different table schemas depending solely on the settings passed to the flattener constructor. This decouples the "what can be extracted" question (answered by the spec) from the "what should be extracted now" question (answered by the hook).

Diagram: Condition Hook Decision Flow

Condition Hook Decision Flow

Type: decision flowchart | sim-id: condition-hook-flow
| Library: vis-network
| Status: Specified

Learning Objective: Trace the evaluation path when a ChildSpec has a condition attribute. Bloom Level: Apply Interactions: Step through the flowchart: ChildSpec encountered -> has condition? -> evaluate_condition() -> True (process subtree) or False (skip subtree). Toggle between Tier C and Tier N settings to see the path change.

Common Condition Hook Patterns

Condition hooks appear in several recurring patterns across TreeSpec deployments:

Feature flags -- The simplest pattern checks a boolean in settings:

@hooks.condition("include_experimental")
def include_experimental(ctx: TreeContext) -> bool:
    return ctx.settings.get("EXPERIMENTAL_FEATURES", False)

Version gating -- Some fields exist only in newer schema versions:

@hooks.condition("v2_fields")
def v2_fields(ctx: TreeContext) -> bool:
    return ctx.version is not None and ctx.version >= "2.0"

Ancestry filtering -- Include a node only when it appears under a specific parent:

@hooks.condition("under_premium_account")
def under_premium_account(ctx: TreeContext) -> bool:
    return "PremiumAccount" in ctx.parent_chain

Field Transform Hooks

A field transform hook is a function that receives a value and a TreeContext, and returns a (possibly modified) value. Unlike condition hooks, which make binary include-or-exclude decisions, field transforms operate on individual field values during extraction. They let you normalize, redact, convert, or enrich data without changing the source objects or the spec.

A field transform is registered with the field_transform() decorator and invoked via apply_transform(). The function signature always takes the raw value as its first argument and the context as its second.

The following example registers a transform that uppercases string values. This might be used to normalize a name field before it reaches the output table.

hooks = HookRegistry()

@hooks.field_transform("upper")
def upper(value, ctx: TreeContext):
    return value.upper() if isinstance(value, str) else value

ctx = TreeContext()
result = hooks.apply_transform("upper", "hello world", ctx)
print(result)  # "HELLO WORLD"

Field transforms are attached to individual fields in the spec via a transform attribute on the field definition. When the flattener extracts that field, it passes the raw value through the named transform before storing it in the output row.

Practical Transform Examples

Date normalization -- Convert various date formats to ISO 8601:

from datetime import datetime

@hooks.field_transform("normalize_date")
def normalize_date(value, ctx: TreeContext):
    if isinstance(value, datetime):
        return value.isoformat()
    if isinstance(value, str):
        for fmt in ("%d/%m/%Y", "%m-%d-%Y", "%Y%m%d"):
            try:
                return datetime.strptime(value, fmt).date().isoformat()
            except ValueError:
                continue
    return value

PII redaction -- Mask sensitive fields in non-production environments:

@hooks.field_transform("redact_email")
def redact_email(value, ctx: TreeContext):
    if ctx.settings.get("ENV") == "production":
        return value
    if isinstance(value, str) and "@" in value:
        local, domain = value.split("@", 1)
        return f"{local[0]}***@{domain}"
    return value

Hook Type	Input	Output	Granularity	Spec Attribute
Condition	TreeContext	bool	Entire subtree	condition on ChildSpec
Field Transform	value + TreeContext	transformed value	Single field	transform on field

Pre-Flatten Hooks

A pre-flatten hook is a function that runs before the tree walk begins. It receives the tree object (or batch of objects) and the context, and can perform any preparatory work: validating data, injecting default values, logging, or modifying the object graph before the flattener traverses it.

Pre-flatten hooks are registered with the pre_flatten() decorator. Unlike condition and transform hooks, which the engines invoke automatically based on spec attributes, pre-flatten hooks are invoked explicitly by the caller or by batch orchestration code.

hooks = HookRegistry()

@hooks.pre_flatten("validate_required_fields")
def validate_required_fields(tree, ctx: TreeContext):
    """Ensure critical fields are present before flattening."""
    if not getattr(tree, "account_id", None):
        raise ValueError("account_id is required before flattening")
    return tree

@hooks.pre_flatten("inject_defaults")
def inject_defaults(tree, ctx: TreeContext):
    """Fill missing optional fields with sensible defaults."""
    if getattr(tree, "currency", None) is None:
        tree.currency = ctx.settings.get("DEFAULT_CURRENCY", "AUD")
    return tree

In the code above, two pre-flatten hooks are registered. The first validates that a required field exists and raises an error if it does not. The second fills in a default currency based on a runtime setting. These hooks run in the order they are invoked by the orchestration layer, giving you fine-grained control over the preparation pipeline.

Pre-flatten hooks are especially valuable in batch processing scenarios where hundreds of source files flow through the same flattener. A pre-flatten hook can normalize differences between files -- adjusting field names, patching missing attributes, or converting legacy formats -- so that the core flattening logic sees a uniform input regardless of source.

Diagram: Pre-Flatten and Post-Build Hook Lifecycle

Pre-Flatten and Post-Build Hook Lifecycle

Type: sequence diagram | sim-id: hook-lifecycle
| Library: vis-network
| Status: Specified

Learning Objective: Understand the execution order of pre-flatten hooks, the core flatten/build operation, and post-build hooks. Bloom Level: Analyze Interactions: Step through the lifecycle: pre-flatten hooks execute in registration order, then the flatten engine runs its tree walk (with condition and transform hooks firing inline), then post-build hooks execute after reconstruction. Click each phase to see which hook types are active.

Post-Build Hooks

A post-build hook is the mirror image of a pre-flatten hook, but on the reconstruction side. It runs after the TreeBuilder has assembled a tree object from table rows. Post-build hooks receive the reconstructed object and the context, and can perform validation, enrichment, or cleanup on the output.

Post-build hooks are registered with the post_build() decorator. Like pre-flatten hooks, they are invoked explicitly rather than being triggered by spec attributes.

hooks = HookRegistry()

@hooks.post_build("compute_derived_fields")
def compute_derived_fields(tree, ctx: TreeContext):
    """Add computed fields that don't exist in the source tables."""
    if hasattr(tree, "subtotal") and hasattr(tree, "tax"):
        tree.total = tree.subtotal + tree.tax
    return tree

@hooks.post_build("audit_log")
def audit_log(tree, ctx: TreeContext):
    """Record that this tree was reconstructed."""
    import logging
    logger = logging.getLogger("treespec.audit")
    logger.info(
        "Rebuilt %s from source=%s",
        type(tree).__name__,
        ctx.source_file or "unknown",
    )
    return tree

The first hook above computes a derived total field from subtotal and tax after the builder has assembled those fields from table rows. The second hook logs an audit trail entry for every reconstructed tree. Both hooks return the tree object (potentially modified) so that downstream consumers receive the enriched version.

Post-build hooks are particularly useful for:

• Derived field computation -- calculating values that span multiple table rows or child nodes.
• Referential integrity checks -- verifying that foreign key relationships are consistent after reconstruction.
• Format conversion -- transforming the rebuilt object into a different representation (JSON, protocol buffer, domain model) before it leaves the pipeline.

Combining Hook Types: A Complete Pipeline

The real power of the hooks system emerges when you combine multiple hook types into a single registry. Each type handles a different concern at a different point in the pipeline, and together they form a processing chain that is fully configurable at runtime.

The following example assembles a complete registry for a financial data pipeline that must handle tiered customers, normalize dates, validate input, and compute totals after reconstruction.

from mountainash_treespec.core.hooks.registry import HookRegistry
from mountainash_treespec.core.hooks.context import TreeContext

hooks = HookRegistry()

# --- Pre-flatten: validate and normalize input ---
@hooks.pre_flatten("validate_account")
def validate_account(tree, ctx: TreeContext):
    if not getattr(tree, "account_id", None):
        raise ValueError("Missing account_id")
    return tree

# --- Condition: tier-based inclusion ---
@hooks.condition("tier_c")
def tier_c(ctx: TreeContext) -> bool:
    return ctx.settings.get("BATCH_TIER") == "C"

@hooks.condition("include_audit_fields")
def include_audit_fields(ctx: TreeContext) -> bool:
    return ctx.settings.get("ENV") == "production"

# --- Field transform: data normalization ---
@hooks.field_transform("normalize_name")
def normalize_name(value, ctx: TreeContext):
    return value.strip().title() if isinstance(value, str) else value

@hooks.field_transform("redact_ssn")
def redact_ssn(value, ctx: TreeContext):
    if ctx.settings.get("ENV") != "production":
        return "***-**-****" if value else value
    return value

# --- Post-build: enrich reconstructed objects ---
@hooks.post_build("compute_totals")
def compute_totals(tree, ctx: TreeContext):
    if hasattr(tree, "line_items"):
        tree.total = sum(item.amount for item in tree.line_items)
    return tree

# Verify all six hooks are registered
print(hooks.hook_names())
# {"validate_account", "tier_c", "include_audit_fields",
#  "normalize_name", "redact_ssn", "compute_totals"}

This registry contains six hooks spanning all four types. A single TreeFlattener and TreeBuilder pair, both constructed with this registry, will validate input, conditionally include credit details, normalize names, redact SSNs in non-production environments, and compute totals after reconstruction -- all without a single line of that logic appearing in the spec file.

Diagram: Complete Hook Pipeline

Complete Hook Pipeline

Type: pipeline diagram | sim-id: complete-hook-pipeline
| Library: vis-network
| Status: Specified

Learning Objective: See how all four hook types compose into a single processing pipeline. Bloom Level: Analyze Interactions: Drag to rearrange the pipeline stages. Click each hook node to see its type signature and when it fires. Toggle "parallel mode" to see the picklability constraint highlighted on each hook.

Table Hooks: Post-Flatten DataFrame Processing

Beyond the four core hook types in HookRegistry, TreeSpec also provides a TableHook mechanism for post-flatten DataFrame processing. A TableHook pairs an optional TypeSpec (for schema conformance) with an optional transform function that operates on Polars DataFrames. This is useful when you need to cast columns to specific types, rename fields, or apply row-level filtering after the flattener has produced its raw output.

The apply_table_hooks() function iterates over the flattened tables, and for each table that has a matching TableHook entry, it first conforms the DataFrame to the TypeSpec schema (casting types, renaming columns) and then applies the transform function. Tables without a matching hook pass through unchanged.

from mountainash_treespec.core.hooks.table_hooks import TableHook

# Define a hook that drops rows with null account_id
# and casts amount to Float64
def clean_accounts(df):
    return df.drop_nulls(subset=["account_id"])

table_hooks = {
    "accounts": TableHook(
        typespec=my_typespec,       # optional: enforce column types
        transform=clean_accounts,   # optional: post-conform transform
    ),
}

TableHooks are validated before execution via validate_table_hooks(), which checks that every key in the hooks dictionary corresponds to a table declared in the TreeSpec. Unknown keys raise a ValueError immediately, catching configuration errors early.

Error Handling

The hooks system raises specific exceptions to help you diagnose problems quickly:

• KeyError -- raised by evaluate_condition() or apply_transform() when the named hook does not exist in the registry. This typically means the spec references a hook name that was never registered.
• UnpicklableHookError -- raised by validate_parallel() when a hook cannot be serialized for multi-process execution.
• TableHookError -- raised by apply_table_hooks() when a conform or transform step fails, wrapping the original exception with the table name and phase (ingress, conform, or transform) for easier debugging.

When a KeyError occurs, the error message includes the exact hook name that was not found, making it straightforward to trace back to the spec attribute or registration code that is mismatched.

Key Takeaways

• The HookRegistry class is the central container for all hooks, organized into four dictionaries by hook type, and uses the Decorator Pattern for registration.
• Condition hooks return a boolean from a TreeContext and control whether entire subtrees are included or excluded during flattening and building.
• Field transform hooks receive a value and a TreeContext and return a transformed value, enabling normalization, redaction, and enrichment at the individual field level.
• Pre-flatten hooks run before the tree walk begins and are used for validation, default injection, and input normalization.
• Post-build hooks run after tree reconstruction and handle derived field computation, audit logging, and format conversion.
• The TreeContext is an immutable dataclass that carries runtime state (settings, version, ancestry, current node) through every hook invocation without risk of cross-contamination.
• All hooks must be picklable (i.e., module-level named functions, not lambdas or closures) when parallel batch processing is used; the validate_parallel() method enforces this constraint.
• Multiple hook types can be composed into a single registry to form a complete processing pipeline that is fully configurable at runtime without modifying the spec.