Source code for numeraire_dataset.zones.steps

"""Deterministic ``frame(s) -> frame`` steps + an open, recipe-hashed step registry.

A **step** is a pure function of its inputs and its explicitly-declared parameters: same inputs +
same params => same output, bit-for-bit. Its identity is ``(name, version, params)``, and its
**recipe hash** is the SHA-256 of canonical JSON that binds every declared input role to its input
recipe or content hash — so the hash *chains*: a clean table's recipe hash transitively pins every
upstream transform and raw vintage. Preprocessing is part of the method, so it is pinned like one.

Steps register in an **open registry** — the same pattern as numeraire's evaluator registry
(``register_step`` / ``get_step`` / ``available_steps``, a module-global dict, ``KeyError`` on a
duplicate unless ``overwrite``). ``@step(...)`` is sugar over :func:`register_step`.
"""

from __future__ import annotations

import datetime as dt
import hashlib
import inspect
import json
import struct
from collections.abc import Callable, Iterable, Mapping
from dataclasses import dataclass
from enum import Enum
from typing import Any
from zoneinfo import ZoneInfo

import numpy as np
import pandas as pd

StepFn = Callable[..., pd.DataFrame]


[docs] @dataclass(frozen=True) class Step: """A registered transform: a pure function plus its identity and input/output artifact names.""" name: str version: int fn: StepFn inputs: tuple[str, ...] output: str
_STEPS: dict[str, Step] = {} _PANDAS_OFFSET_TYPES = frozenset( candidate for candidate in vars(pd.offsets).values() if isinstance(candidate, type) and issubclass(candidate, pd.DateOffset) )
[docs] def register_step(step_obj: Step, *, overwrite: bool = False) -> Step: """Register ``step_obj`` under its ``name``. Raises on a duplicate unless ``overwrite``.""" if not overwrite and step_obj.name in _STEPS: raise KeyError(f"step {step_obj.name!r} already registered") signature = inspect.signature(step_obj.fn) if any( parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in signature.parameters.values() ): raise TypeError( f"step {step_obj.name!r} must declare every parameter explicitly; " "**kwargs is unsupported" ) _STEPS[step_obj.name] = step_obj return step_obj
[docs] def get_step(name: str) -> Step: """Return the step registered under ``name``.""" try: return _STEPS[name] except KeyError: raise KeyError(f"no step registered as {name!r}") from None
[docs] def available_steps() -> tuple[str, ...]: """Return the names of all registered steps, sorted.""" return tuple(sorted(_STEPS))
[docs] def step( *, name: str, version: int, inputs: list[str], output: str, overwrite: bool = False ) -> Callable[[StepFn], StepFn]: """Decorator: register ``fn`` as a step and return it unchanged (still directly callable).""" def decorate(fn: StepFn) -> StepFn: register_step( Step(name=name, version=version, fn=fn, inputs=tuple(inputs), output=output), overwrite=overwrite, ) return fn return decorate
RECIPE_HASH_VERSION = 2 CONTENT_DIGEST_VERSION = 2 def _type_name(value: object) -> str: cls = type(value) return f"{cls.__module__}.{cls.__qualname__}" def _canonical_json(value: object) -> str: return json.dumps( value, sort_keys=True, separators=(",", ":"), ensure_ascii=False, allow_nan=False, ) def _canonical_timezone(value: dt.tzinfo | None) -> dict[str, object] | None: """Canonicalize timezones with stable public identities; reject opaque custom state.""" if value is None: return None if type(value) is dt.timezone: offset = value.utcoffset(None) return { "type": "datetime.timezone", "offset": None if offset is None else _canonical_parameter(offset), "name": value.tzname(None), } if type(value) is ZoneInfo: return {"type": "zoneinfo.ZoneInfo", "key": value.key} zone = getattr(value, "zone", None) if type(zone) is str and ( type(value).__module__ == "pytz" or type(value).__module__.startswith("pytz.") ): return {"type": "pytz", "class": _type_name(value), "zone": zone} raise TypeError( f"unsupported timezone type {_type_name(value)!r}; use datetime.timezone, ZoneInfo, " "or a named pytz zone" ) def _canonical_parameter(value: object) -> dict[str, object]: """Return a type-tagged, JSON-safe parameter value or fail closed. Recipe parameters are executable inputs: a step may legitimately branch on ``list`` versus ``tuple`` or on a timestamp versus a string with the same display. Standard JSON and ``default=str`` erase those distinctions, so every supported scalar/container carries an explicit type tag. Unsupported stateful objects are rejected instead of being stringified into a potentially colliding or process-dependent recipe. """ if value is None: return {"type": "none"} if type(value) is bool: return {"type": "bool", "value": value} if type(value) is int: return {"type": "int", "value": str(value)} if type(value) is float: return {"type": "float", "bytes": struct.pack("!d", value).hex()} if type(value) is str: return {"type": "str", "value": value} if type(value) is bytes: return {"type": "bytes", "value": value.hex()} if type(value) is bytearray: return {"type": "bytearray", "value": bytes(value).hex()} if value is pd.NA: return {"type": "pandas.NA"} if value is pd.NaT: return {"type": "pandas.NaT"} if isinstance(value, Enum): state = { key: item for key, item in vars(value).items() if key not in {"_value_", "_name_", "__objclass__", "_sort_order_"} } return { "type": "enum", "class": _type_name(value), "name": value.name, "value": _canonical_parameter(value.value), "state": _canonical_parameter(state), } if isinstance(value, np.generic): if not type(value).__module__.startswith("numpy"): raise TypeError( f"unsupported custom NumPy scalar type {_type_name(value)!r}; " "convert it to an exact NumPy scalar or built-in value" ) if value.dtype.hasobject: raise TypeError(f"recipe parameter {_type_name(value)!r} contains object state") if value.dtype.fields is not None or value.dtype.subdtype is not None: raise TypeError( f"recipe parameter {_type_name(value)!r} has a structured NumPy dtype; " "convert it to explicit built-in values" ) return { "type": "numpy_scalar", "class": _type_name(value), "dtype": value.dtype.str, "bytes": value.tobytes().hex(), } if type(value) is pd.Timestamp: scalar = value.asm8 return { "type": "pandas.Timestamp", "value": str(scalar.view("i8").item()), "unit": scalar.dtype.str, "timezone": _canonical_timezone(value.tzinfo), "fold": value.fold, } if type(value) is pd.Timedelta: scalar = value.asm8 return { "type": "pandas.Timedelta", "value": str(scalar.view("i8").item()), "unit": scalar.dtype.str, } if type(value) is pd.Period: return { "type": "pandas.Period", "ordinal": str(value.ordinal), "frequency": value.freqstr, } if type(value) is pd.Interval: return { "type": "pandas.Interval", "left": _canonical_parameter(value.left), "right": _canonical_parameter(value.right), "closed": value.closed, } if isinstance(value, pd.DateOffset): if type(value) not in _PANDAS_OFFSET_TYPES: raise TypeError( f"unsupported recipe parameter type {_type_name(value)!r}; " "custom pandas offsets must be converted to explicit built-in values" ) offset_params = getattr(value, "_params", None) if type(offset_params) is not tuple: raise TypeError(f"pandas offset {_type_name(value)!r} has no canonical parameter tuple") return { "type": "pandas.DateOffset", "class": _type_name(value), "params": _canonical_parameter(offset_params), } if type(value) is dt.datetime: return { "type": "datetime.datetime", "value": value.isoformat(), "fold": value.fold, "timezone": _canonical_timezone(value.tzinfo), } if type(value) is dt.date: return {"type": "datetime.date", "value": value.isoformat()} if type(value) is dt.time: return { "type": "datetime.time", "value": value.isoformat(), "fold": value.fold, "timezone": _canonical_timezone(value.tzinfo), } if type(value) is dt.timedelta: return { "type": "datetime.timedelta", "days": str(value.days), "seconds": str(value.seconds), "microseconds": str(value.microseconds), } if type(value) is dict: items = [ [_canonical_parameter(key), _canonical_parameter(item)] for key, item in value.items() ] # Nested dict order is executable state in Python (a step may iterate it), unlike the # top-level mapping of named keyword parameters whose order is canonicalized. return {"type": "mapping", "items": items} if type(value) is list: return { "type": "list", "items": [_canonical_parameter(item) for item in value], } if type(value) is tuple: return { "type": "tuple", "items": [_canonical_parameter(item) for item in value], } if isinstance(value, set) and type(value) is set: items = [_canonical_parameter(item) for item in value] items.sort(key=_canonical_json) return {"type": "set", "items": items} if isinstance(value, frozenset) and type(value) is frozenset: items = [_canonical_parameter(item) for item in value] items.sort(key=_canonical_json) return {"type": "frozenset", "items": items} raise TypeError( f"unsupported recipe parameter type {_type_name(value)!r}; use explicit JSON-like, " "datetime/pandas temporal, enum, or NumPy scalar values" ) def _canonical_params(params: Mapping[str, object]) -> dict[str, object]: for key in params: if type(key) is not str: raise TypeError("recipe parameter names must be plain strings") return {key: _canonical_parameter(value) for key, value in params.items()} def _recipe_hash_from_canonical( name: str, version: int, canonical_params: Mapping[str, object], input_hashes: Mapping[str, str], ) -> str: for role, input_hash in input_hashes.items(): if type(role) is not str or type(input_hash) is not str: raise TypeError("recipe input roles and hashes must be plain strings") payload = { "recipe_hash_version": RECIPE_HASH_VERSION, "name": name, "version": version, "params": dict(canonical_params), "inputs": dict(input_hashes), } blob = _canonical_json(payload) return "sha256:" + hashlib.sha256(blob.encode("utf-8")).hexdigest()
[docs] def recipe_hash( name: str, version: int, params: dict[str, Any], input_hashes: Mapping[str, str], ) -> str: """Hash a recipe whose canonical input mapping preserves each input's semantic role. Mapping keys are the step's declared input names. Unlike the legacy value-only list, this representation cannot identify ``left=X, right=Y`` with ``left=Y, right=X``. Contract version 2 is embedded in the payload so every clean recipe intentionally moves from the legacy hash. """ if not isinstance(input_hashes, Mapping): raise TypeError( "input_hashes must be a role-to-hash mapping such as {'source': 'sha256:...'}; " "the legacy hash-list form was removed by recipe-hash contract version 2" ) return _recipe_hash_from_canonical(name, version, _canonical_params(params), input_hashes)
def _label_schema(value: object) -> dict[str, str]: """Represent a column/index label without collapsing e.g. ``1`` and ``"1"``.""" return {"type": _type_name(value), "repr": repr(value)} def _update_type_tag(digest: Any, value: object) -> None: """Hash scalar/container type structure without materializing an object-column copy.""" name = _type_name(value).encode("utf-8") digest.update(len(name).to_bytes(4, "big")) digest.update(name) if isinstance(value, tuple): digest.update(len(value).to_bytes(4, "big")) for item in value: _update_type_tag(digest, item) def _object_type_digest(values: Iterable[object]) -> str: """Distinguish object scalars that pandas' value hash intentionally coalesces.""" digest = hashlib.sha256() for value in values: _update_type_tag(digest, value) return "sha256:" + digest.hexdigest() def _dtype_schema(dtype: object) -> dict[str, Any]: """Return the pandas dtype metadata that can affect downstream frame semantics.""" out: dict[str, Any] = { "type": _type_name(dtype), "text": str(dtype), "repr": repr(dtype), } if isinstance(dtype, pd.CategoricalDtype): categories = dtype.categories category_hash = hashlib.sha256( pd.util.hash_pandas_object(categories, index=True).to_numpy().tobytes() ).hexdigest() out.update( { "ordered": dtype.ordered, "categories_dtype": _dtype_schema(categories.dtype), "categories_hash": f"sha256:{category_hash}", "categories_type_hash": _object_type_digest(categories), } ) if isinstance(dtype, pd.DatetimeTZDtype): out.update({"unit": dtype.unit, "tz": str(dtype.tz)}) return out def _index_schema(index: pd.Index) -> dict[str, Any]: """Return index structure not represented by pandas' row-value hash.""" out: dict[str, Any] = { "type": _type_name(index), "names": [_label_schema(name) for name in index.names], "nlevels": index.nlevels, } if isinstance(index, pd.MultiIndex): levels: list[dict[str, Any]] = [] for level in index.levels: value_hash = hashlib.sha256( pd.util.hash_pandas_object(level, index=True).to_numpy().tobytes() ).hexdigest() levels.append( { "dtype": _dtype_schema(level.dtype), "values_hash": f"sha256:{value_hash}", "value_type_hash": _object_type_digest(level), } ) codes = hashlib.sha256() for code in index.codes: canonical = code.astype("int64", copy=False) codes.update(len(canonical).to_bytes(8, "big")) codes.update(canonical.tobytes()) out["levels"] = levels out["codes_hash"] = "sha256:" + codes.hexdigest() out["sortorder"] = getattr(index, "sortorder", None) out["value_type_hash"] = _object_type_digest(index) else: out["dtype"] = _dtype_schema(index.dtype) if pd.api.types.is_object_dtype(index.dtype): out["value_type_hash"] = _object_type_digest(index) if isinstance(index, pd.RangeIndex): out["range"] = {"start": index.start, "stop": index.stop, "step": index.step} freq = getattr(index, "freqstr", None) if freq is not None: out["frequency"] = str(freq) return out
[docs] def content_digest(frame: pd.DataFrame) -> str: """Digest a frame's values and semantic schema, preserving row/column order. Pandas' value hash deliberately treats several dtype representations as equivalent. A data provenance digest cannot: object versus categorical values, category order, timezone, extension dtype, and index metadata can all change downstream transforms while leaving the displayed values unchanged. The canonical schema payload therefore accompanies the value hash. ``CONTENT_DIGEST_VERSION`` must move whenever this contract changes. """ h = hashlib.sha256() schema = { "version": CONTENT_DIGEST_VERSION, "columns": _index_schema(frame.columns), "column_dtypes": [_dtype_schema(dtype) for dtype in frame.dtypes], "object_value_type_hashes": [ _object_type_digest(frame.iloc[:, position]) if pd.api.types.is_object_dtype(dtype) else None for position, dtype in enumerate(frame.dtypes) ], "index": _index_schema(frame.index), } h.update(json.dumps(schema, sort_keys=True, separators=(",", ":")).encode("utf-8")) h.update(pd.util.hash_pandas_object(frame.columns, index=True).to_numpy().tobytes()) h.update(pd.util.hash_pandas_object(frame, index=True).to_numpy().tobytes()) return "sha256:" + h.hexdigest()
[docs] @dataclass(frozen=True) class Built: """A built frame plus the step identity and explicit inputs committed by its recipe hash. Obtain instances from :func:`run_step`; its metadata fields form one validated recipe snapshot and are not a stable direct-construction API. """ name: str frame: pd.DataFrame recipe_hash: str content_digest: str canonical_params: dict[str, object] input_hashes: dict[str, str] step_name: str step_version: int recipe_hash_version: int = RECIPE_HASH_VERSION
[docs] def run_step( name: str, frames: dict[str, pd.DataFrame], *, params: dict[str, Any] | None = None, input_hashes: dict[str, str] | None = None, ) -> Built: """Run a registered step on named input frames, returning the output + its provenance hashes. ``frames`` maps each declared input name to its frame; ``input_hashes`` maps them to the recipe (or raw-content) hash that identifies them, so the output's :func:`recipe_hash` chains onto its inputs. ``params`` are the step's explicit parameters (they participate in the recipe hash). """ step_obj = get_step(name) params = dict(params or {}) input_hashes = dict(input_hashes or {}) missing = [i for i in step_obj.inputs if i not in frames] if missing: raise ValueError(f"step {name!r} missing input frame(s) {missing}") declared_inputs = set(step_obj.inputs) extra_frames = sorted(set(frames) - declared_inputs) if extra_frames: raise ValueError(f"step {name!r} received unexpected input frame role(s) {extra_frames}") extra_hashes = sorted(set(input_hashes) - declared_inputs) if extra_hashes: raise ValueError(f"step {name!r} received unexpected input-hash role(s) {extra_hashes}") ordered = [frames[i] for i in step_obj.inputs] canonical_params = _canonical_params(params) resolved_hashes = { input_name: ( input_hashes[input_name] if input_name in input_hashes else content_digest(frames[input_name]) ) for input_name in step_obj.inputs } # Freeze the complete recipe at call entry. A step is contractually pure, but recording a # digest or mutable parameter after accidental in-place mutation would pin the wrong input. rhash = _recipe_hash_from_canonical( step_obj.name, step_obj.version, canonical_params, resolved_hashes ) out = step_obj.fn(*ordered, **params) return Built( name=step_obj.output, frame=out, recipe_hash=rhash, content_digest=content_digest(out), canonical_params=canonical_params, input_hashes=resolved_hashes, step_name=step_obj.name, step_version=step_obj.version, )