From e2f7373b2a915992ae303f168d63b49a4ff06c05 Mon Sep 17 00:00:00 2001 From: Dylan Mordaunt <15080672+edithatogo@users.noreply.github.com> Date: Thu, 9 Jul 2026 11:24:46 +1000 Subject: [PATCH 1/3] Add versioned computation trace export (policyengine.trace.v1) Expose FullTracer.to_trace() and Simulation.to_trace() over the existing FlatTrace serialization so household calculations can export a stable, JSON-compatible audit document without changing formula evaluation. Closes PolicyEngine/policyengine-core#512 --- changelog.d/versioned-trace-export.added.md | 1 + policyengine_core/simulations/simulation.py | 46 +++++++ policyengine_core/tracers/full_tracer.py | 129 +++++++++++++++++++- tests/core/test_tracers.py | 61 +++++++++ 4 files changed, 236 insertions(+), 1 deletion(-) create mode 100644 changelog.d/versioned-trace-export.added.md diff --git a/changelog.d/versioned-trace-export.added.md b/changelog.d/versioned-trace-export.added.md new file mode 100644 index 00000000..77643ee9 --- /dev/null +++ b/changelog.d/versioned-trace-export.added.md @@ -0,0 +1 @@ +Added `FullTracer.to_trace()` and `Simulation.to_trace()` for versioned computation trace export (`policyengine.trace.v1`). diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index b7f52591..04d554ca 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -528,6 +528,52 @@ def trace(self, trace: SimpleTracer) -> None: else: self.tracer = SimpleTracer() + def to_trace( + self, + format: str = "policyengine.trace.v1", + *, + model: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Export a versioned computation trace for this simulation. + + Requires ``trace=True`` (or assigning ``simulation.trace = True``) + so a :class:`~policyengine_core.tracers.FullTracer` is installed. + + :param format: Trace document format identifier. + :param model: Optional country-package metadata override. When omitted, + a best-effort package name/version is taken from the tax-benefit + system module when available. + :returns: JSON-compatible trace document (see + :meth:`~policyengine_core.tracers.FullTracer.to_trace`). + """ + if not isinstance(self.tracer, FullTracer): + raise ValueError( + "Computation traces require FullTracer. " + "Create the simulation with trace=True, or set simulation.trace = True " + "before calculating, then call to_trace()." + ) + if model is None: + model = self._model_metadata_for_trace() + return self.tracer.to_trace(format=format, model=model) + + def _model_metadata_for_trace(self) -> Optional[Dict[str, Any]]: + tax_benefit_system = getattr(self, "tax_benefit_system", None) + if tax_benefit_system is None: + return None + module = type(tax_benefit_system).__module__ + package_name = module.split(".", 1)[0] if module else None + if not package_name or package_name == "policyengine_core": + return None + metadata: Dict[str, Any] = {"package": package_name.replace("_", "-")} + try: + import importlib.metadata + + distribution_name = package_name.replace("_", "-") + metadata["version"] = importlib.metadata.version(distribution_name) + except Exception: + pass + return metadata + def link_to_entities_instances(self) -> None: for _key, entity_instance in self.populations.items(): entity_instance.simulation = self diff --git a/policyengine_core/tracers/full_tracer.py b/policyengine_core/tracers/full_tracer.py index 3cab802e..9b56302c 100644 --- a/policyengine_core/tracers/full_tracer.py +++ b/policyengine_core/tracers/full_tracer.py @@ -1,8 +1,9 @@ from __future__ import annotations +import re import time import typing -from typing import Dict, Iterator, List, Optional, Union +from typing import Any, Dict, Iterator, List, Optional, Union from .. import tracers @@ -14,6 +15,25 @@ Stack = List[Dict[str, Union[str, Period]]] +TRACE_FORMAT_V1 = "policyengine.trace.v1" +# Matches FlatTrace.key(): ``name`` +_TRACE_KEY_RE = re.compile( + r"^(?P.+)<(?P[^<>]+), \((?P[^()]*)\)>$" +) + + +def parse_trace_key(key: str) -> Dict[str, str]: + """Parse a serialized flat-trace node key into name/period/branch.""" + match = _TRACE_KEY_RE.match(key) + if not match: + raise ValueError(f"Invalid PolicyEngine trace key: {key!r}") + return { + "name": match.group("name"), + "period": match.group("period").strip(), + "branch": match.group("branch").strip(), + } + + class FullTracer: _simple_tracer: tracers.SimpleTracer _trees: list @@ -172,3 +192,110 @@ def _browse_node(node): for node in self._trees: yield from _browse_node(node) + + def to_trace( + self, + format: str = TRACE_FORMAT_V1, + *, + model: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Export a versioned, JSON-compatible computation trace. + + Builds on :meth:`get_serialized_flat_trace` so existing tracer data + (nodes, dependencies, parameters, values, timings) is preserved under a + stable document shape suitable for audit, debugging, and downstream + projection without requiring PolicyEngine to adopt an external schema. + + :param format: Trace document format identifier. Currently only + ``policyengine.trace.v1`` is supported. + :param model: Optional country-package / model metadata + (for example ``{"package": "policyengine-us", "version": "…"}``). + :returns: A plain ``dict`` ready for ``json.dumps``. + :raises ValueError: If ``format`` is not a supported version string. + """ + if format != TRACE_FORMAT_V1: + raise ValueError( + f"Unsupported trace format {format!r}. " + f"Supported formats: {TRACE_FORMAT_V1!r}" + ) + + engine = self._engine_metadata() + flat = self.get_serialized_flat_trace() + nodes: List[Dict[str, Any]] = [] + parameters: Dict[str, Any] = {} + + for key, payload in flat.items(): + parsed = parse_trace_key(key) + node: Dict[str, Any] = { + "id": key, + "variable": parsed["name"], + "period": parsed["period"], + "branch": parsed["branch"], + "dependencies": list(payload.get("dependencies", [])), + "parameters": dict(payload.get("parameters", {})), + "value": payload.get("value"), + } + if "calculation_time" in payload: + node["calculation_time"] = payload["calculation_time"] + if "formula_time" in payload: + node["formula_time"] = payload["formula_time"] + nodes.append(node) + + for parameter_key, parameter_value in payload.get("parameters", {}).items(): + if parameter_key in parameters: + continue + try: + parameter_parsed = parse_trace_key(parameter_key) + except ValueError: + parameters[parameter_key] = { + "id": parameter_key, + "value": parameter_value, + } + continue + parameters[parameter_key] = { + "id": parameter_key, + "name": parameter_parsed["name"], + "instant": parameter_parsed["period"], + "branch": parameter_parsed["branch"], + "value": parameter_value, + } + + document: Dict[str, Any] = { + "format": format, + "engine": engine, + "calculation": { + "roots": [ + { + "id": self.flat_trace.key(tree), + "variable": tree.name, + "period": str(tree.period), + "branch": tree.branch_name, + } + for tree in self._trees + ], + }, + "nodes": nodes, + "parameters": parameters, + } + if model is not None: + document["model"] = model + return document + + @staticmethod + def _engine_metadata() -> Dict[str, Any]: + try: + from policyengine_core.build_metadata import get_runtime_metadata + + metadata = get_runtime_metadata() + engine: Dict[str, Any] = { + "package": metadata.get("name", "policyengine-core"), + "version": metadata.get("version", "unknown"), + } + if metadata.get("git_sha"): + engine["git_sha"] = metadata["git_sha"] + return engine + except Exception: + return { + "package": "policyengine-core", + "version": "unknown", + } diff --git a/tests/core/test_tracers.py b/tests/core/test_tracers.py index e653f39b..8406c6ed 100644 --- a/tests/core/test_tracers.py +++ b/tests/core/test_tracers.py @@ -529,3 +529,64 @@ def test_browse_trace(): browsed_nodes = [node.name for node in tracer.browse_trace()] assert browsed_nodes == ["B", "C", "D", "E", "F"] + + +def test_to_trace_versioned_export(): + tracer = FullTracer() + tracer.record_calculation_start("income_tax", 2017) + tracer.record_calculation_start("salary", 2017) + tracer.record_calculation_result(np.asarray([2000])) + tracer.record_calculation_end() + tracer.record_parameter_access("taxes.rate", "2017-01-01", "default", 0.15) + tracer.record_calculation_result(np.asarray([300])) + tracer.record_calculation_end() + + trace = tracer.to_trace() + + assert trace["format"] == "policyengine.trace.v1" + assert trace["engine"]["package"] == "policyengine-core" + assert "version" in trace["engine"] + assert trace["calculation"]["roots"] == [ + { + "id": "income_tax<2017, (default)>", + "variable": "income_tax", + "period": "2017", + "branch": "default", + } + ] + + nodes_by_id = {node["id"]: node for node in trace["nodes"]} + root = nodes_by_id["income_tax<2017, (default)>"] + assert root["variable"] == "income_tax" + assert root["period"] == "2017" + assert root["branch"] == "default" + assert root["dependencies"] == ["salary<2017, (default)>"] + assert root["value"] == [300] + assert "taxes.rate<2017-01-01, (default)>" in root["parameters"] + assert root["parameters"]["taxes.rate<2017-01-01, (default)>"] == approx(0.15) + assert "calculation_time" in root + assert "formula_time" in root + + assert nodes_by_id["salary<2017, (default)>"]["value"] == [2000] + assert "taxes.rate<2017-01-01, (default)>" in trace["parameters"] + assert ( + trace["parameters"]["taxes.rate<2017-01-01, (default)>"]["name"] == "taxes.rate" + ) + + +def test_to_trace_rejects_unknown_format(): + tracer = FullTracer() + tracer.record_calculation_start("salary", 2017) + tracer.record_calculation_end() + + with raises(ValueError, match="Unsupported trace format"): + tracer.to_trace(format="policyengine.trace.v0") + + +def test_to_trace_includes_optional_model_metadata(): + tracer = FullTracer() + tracer.record_calculation_start("salary", 2017) + tracer.record_calculation_end() + + trace = tracer.to_trace(model={"package": "policyengine-us", "version": "0.0.0"}) + assert trace["model"] == {"package": "policyengine-us", "version": "0.0.0"} From 7bedc2ca039270c9ff3f888c54e93687cc466d93 Mon Sep 17 00:00:00 2001 From: Dylan Mordaunt <15080672+edithatogo@users.noreply.github.com> Date: Sat, 11 Jul 2026 17:50:40 +1000 Subject: [PATCH 2/3] test: cover Simulation.to_trace and unparseable parameter keys Raise patch coverage for the versioned trace export helpers that Codecov flagged on PR #515. Co-authored-by: Cursor --- tests/core/test_tracers.py | 77 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 77 insertions(+) diff --git a/tests/core/test_tracers.py b/tests/core/test_tracers.py index 8406c6ed..37e25add 100644 --- a/tests/core/test_tracers.py +++ b/tests/core/test_tracers.py @@ -590,3 +590,80 @@ def test_to_trace_includes_optional_model_metadata(): trace = tracer.to_trace(model={"package": "policyengine-us", "version": "0.0.0"}) assert trace["model"] == {"package": "policyengine-us", "version": "0.0.0"} + + +def test_simulation_to_trace_requires_full_tracer(): + simulation = StubSimulation() + simulation.tracer = SimpleTracer() + + with raises(ValueError, match="FullTracer"): + simulation.to_trace() + + +def test_simulation_to_trace_exports_full_tracer_document(): + simulation = StubSimulation() + simulation.tracer = FullTracer() + simulation.tax_benefit_system = None + simulation.tracer.record_calculation_start("salary", 2017) + simulation.tracer.record_calculation_result(np.asarray([1000])) + simulation.tracer.record_calculation_end() + + trace = simulation.to_trace(model={"package": "policyengine-us", "version": "1.0.0"}) + + assert trace["format"] == "policyengine.trace.v1" + assert trace["model"] == {"package": "policyengine-us", "version": "1.0.0"} + assert trace["calculation"]["roots"][0]["variable"] == "salary" + assert trace["nodes"][0]["value"] == [1000] + + +def test_simulation_to_trace_infers_model_metadata_from_tax_benefit_system(): + class FakeCountrySystem: + pass + + FakeCountrySystem.__module__ = "policyengine_us.system" + + simulation = StubSimulation() + simulation.tracer = FullTracer() + simulation.tax_benefit_system = FakeCountrySystem() + simulation.tracer.record_calculation_start("salary", 2017) + simulation.tracer.record_calculation_end() + + trace = simulation.to_trace() + + assert trace["model"]["package"] == "policyengine-us" + + +def test_simulation_to_trace_skips_core_package_as_model_metadata(): + class CoreLikeSystem: + pass + + CoreLikeSystem.__module__ = "policyengine_core.country_template" + + simulation = StubSimulation() + simulation.tracer = FullTracer() + simulation.tax_benefit_system = CoreLikeSystem() + simulation.tracer.record_calculation_start("salary", 2017) + simulation.tracer.record_calculation_end() + + trace = simulation.to_trace() + assert "model" not in trace + + +def test_to_trace_records_unparseable_parameter_keys(): + tracer = FullTracer() + tracer.record_calculation_start("salary", 2017) + # Inject a flat-trace parameter key that parse_trace_key cannot parse. + tracer.record_calculation_end() + flat = tracer.get_serialized_flat_trace() + key = next(iter(flat)) + flat[key]["parameters"] = {"not-a-valid-trace-key": 0.1} + # Monkeypatch get_serialized_flat_trace for this call path via temporary override + original = tracer.get_serialized_flat_trace + tracer.get_serialized_flat_trace = lambda: flat + try: + trace = tracer.to_trace() + finally: + tracer.get_serialized_flat_trace = original + + assert "not-a-valid-trace-key" in trace["parameters"] + assert trace["parameters"]["not-a-valid-trace-key"]["value"] == 0.1 From b9979f40ddc9439a0c82947cd998eb32a78f1b14 Mon Sep 17 00:00:00 2001 From: Dylan Mordaunt <15080672+edithatogo@users.noreply.github.com> Date: Sat, 11 Jul 2026 18:04:37 +1000 Subject: [PATCH 3/3] refactor(trace): build to_trace nodes from TraceNode fields Address review feedback: avoid regex-parsing flat keys for node identity, document v1-absent fields and TRO complementarity, keep parse_trace_key for flat-key callers only. Co-authored-by: Cursor --- policyengine_core/tracers/full_tracer.py | 92 ++++++++++++++---------- tests/core/test_tracers.py | 42 ++++++----- 2 files changed, 82 insertions(+), 52 deletions(-) diff --git a/policyengine_core/tracers/full_tracer.py b/policyengine_core/tracers/full_tracer.py index 9b56302c..4afb10cb 100644 --- a/policyengine_core/tracers/full_tracer.py +++ b/policyengine_core/tracers/full_tracer.py @@ -201,10 +201,26 @@ def to_trace( ) -> Dict[str, Any]: """Export a versioned, JSON-compatible computation trace. - Builds on :meth:`get_serialized_flat_trace` so existing tracer data - (nodes, dependencies, parameters, values, timings) is preserved under a - stable document shape suitable for audit, debugging, and downstream - projection without requiring PolicyEngine to adopt an external schema. + Node identity fields (``variable`` / ``period`` / ``branch``) are taken + from structured :class:`~policyengine_core.tracers.TraceNode` attributes + via :meth:`browse_trace`, not by regex-parsing serialized flat-trace + keys. :func:`parse_trace_key` remains available for callers that only + have a flat key string (for example parameter maps without a node). + + **v1 intentionally omits** entity/count on ``calculation`` roots and + source/version on ``parameters``. Those were listed as optional + "minimum useful fields" in the provenance issue and can land in a later + format version without breaking ``policyengine.trace.v1`` consumers. + + This document is *per-computation* provenance. It is complementary to + PolicyEngine's TRACE Transparent Research Object (TRO) surface in + ``policyengine`` (JSON-LD under ``https://policyengine.org/trace/0.1#``), + which pins release/composition artifacts. A future TRO may embed or + reference a ``policyengine.trace.v1`` payload as its runtime trace; the + format id and TRO namespace are intentionally distinct. + + Values are full serialized vectors (household-audit friendly). Optional + summarization for population-scale traces is deferred. :param format: Trace document format identifier. Currently only ``policyengine.trace.v1`` is supported. @@ -220,44 +236,48 @@ def to_trace( ) engine = self._engine_metadata() - flat = self.get_serialized_flat_trace() + flat_trace = self.flat_trace nodes: List[Dict[str, Any]] = [] parameters: Dict[str, Any] = {} - - for key, payload in flat.items(): - parsed = parse_trace_key(key) + seen_node_ids: set[str] = set() + + for tree_node in self.browse_trace(): + node_id = flat_trace.key(tree_node) + # Cache reads can revisit the same calculation; keep the first + # structured record (same non-overwriting rule as FlatTrace). + if node_id in seen_node_ids: + continue + seen_node_ids.add(node_id) + + parameter_map = { + flat_trace.key(parameter): flat_trace.serialize(parameter.value) + for parameter in tree_node.parameters + } node: Dict[str, Any] = { - "id": key, - "variable": parsed["name"], - "period": parsed["period"], - "branch": parsed["branch"], - "dependencies": list(payload.get("dependencies", [])), - "parameters": dict(payload.get("parameters", {})), - "value": payload.get("value"), + "id": node_id, + "variable": tree_node.name, + "period": str(tree_node.period), + "branch": tree_node.branch_name, + "dependencies": [ + flat_trace.key(child) for child in tree_node.children + ], + "parameters": parameter_map, + "value": flat_trace.serialize(tree_node.value), + "calculation_time": tree_node.calculation_time(), + "formula_time": tree_node.formula_time(), } - if "calculation_time" in payload: - node["calculation_time"] = payload["calculation_time"] - if "formula_time" in payload: - node["formula_time"] = payload["formula_time"] nodes.append(node) - for parameter_key, parameter_value in payload.get("parameters", {}).items(): - if parameter_key in parameters: - continue - try: - parameter_parsed = parse_trace_key(parameter_key) - except ValueError: - parameters[parameter_key] = { - "id": parameter_key, - "value": parameter_value, - } + for parameter in tree_node.parameters: + parameter_id = flat_trace.key(parameter) + if parameter_id in parameters: continue - parameters[parameter_key] = { - "id": parameter_key, - "name": parameter_parsed["name"], - "instant": parameter_parsed["period"], - "branch": parameter_parsed["branch"], - "value": parameter_value, + parameters[parameter_id] = { + "id": parameter_id, + "name": parameter.name, + "instant": str(parameter.period), + "branch": parameter.branch_name, + "value": flat_trace.serialize(parameter.value), } document: Dict[str, Any] = { @@ -266,7 +286,7 @@ def to_trace( "calculation": { "roots": [ { - "id": self.flat_trace.key(tree), + "id": flat_trace.key(tree), "variable": tree.name, "period": str(tree.period), "branch": tree.branch_name, diff --git a/tests/core/test_tracers.py b/tests/core/test_tracers.py index 37e25add..67430cfd 100644 --- a/tests/core/test_tracers.py +++ b/tests/core/test_tracers.py @@ -17,6 +17,7 @@ TraceNode, TracingParameterNodeAtInstant, ) +from policyengine_core.tracers.full_tracer import parse_trace_key from .parameters_fancy_indexing.test_fancy_indexing import parameters @@ -649,21 +650,30 @@ class CoreLikeSystem: assert "model" not in trace -def test_to_trace_records_unparseable_parameter_keys(): +def test_to_trace_reads_structured_fields_from_trace_nodes(): + """Nodes/parameters come from TraceNode attributes, not key regex parsing.""" tracer = FullTracer() - tracer.record_calculation_start("salary", 2017) - # Inject a flat-trace parameter key that parse_trace_key cannot parse. + tracer.record_calculation_start("income_tax", 2017, branch_name="reform") + tracer.record_parameter_access("taxes.rate", "2017-01-01", "reform", 0.2) + tracer.record_calculation_result(np.asarray([100])) tracer.record_calculation_end() - flat = tracer.get_serialized_flat_trace() - key = next(iter(flat)) - flat[key]["parameters"] = {"not-a-valid-trace-key": 0.1} - # Monkeypatch get_serialized_flat_trace for this call path via temporary override - original = tracer.get_serialized_flat_trace - tracer.get_serialized_flat_trace = lambda: flat - try: - trace = tracer.to_trace() - finally: - tracer.get_serialized_flat_trace = original - - assert "not-a-valid-trace-key" in trace["parameters"] - assert trace["parameters"]["not-a-valid-trace-key"]["value"] == 0.1 + + trace = tracer.to_trace() + root = trace["nodes"][0] + assert root["variable"] == "income_tax" + assert root["period"] == "2017" + assert root["branch"] == "reform" + + parameter = trace["parameters"]["taxes.rate<2017-01-01, (reform)>"] + assert parameter["name"] == "taxes.rate" + assert parameter["instant"] == "2017-01-01" + assert parameter["branch"] == "reform" + assert parameter["value"] == approx(0.2) + + +def test_parse_trace_key_still_available_for_flat_keys(): + parsed = parse_trace_key("salary<2017, (default)>") + assert parsed == {"name": "salary", "period": "2017", "branch": "default"} + + with raises(ValueError, match="Invalid PolicyEngine trace key"): + parse_trace_key("not-a-valid-trace-key")