diff --git a/AGENTS.md b/AGENTS.md
index b980d9a53..88818333e 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -29,6 +29,10 @@ When changing correlation identifiers, telemetry transport, spans, stage
names, logging, or observability failure handling, read
`docs/engineering/skills/observability.md`.
+When changing code that aggregates simulation microdata (`MicroSeries`,
+`MicroDataFrame` or economy result columns), read
+`docs/engineering/skills/microdata-weights.md`.
+
When modifying the `Makefile`, read
`docs/engineering/skills/repository-maintenance.md`.
@@ -94,6 +98,11 @@ To add or change a rule, read `docs/engineering/skills/rules-registry.md`.
- Checked: `api.migration.alembic-v1-qualified`, `api.migration.alembic-v2-qualified`, `api.migration.bundle-supported-by-simulation-api`, `api.migration.contracts-consistent`, `api.migration.cross-database-integration`
- `api.migration.alembic-autogenerated-only`: Alembic revisions are autogenerated against an explicitly selected v1 or v2 configuration; never hand-written DDL or data operations. (doc-only, not checked: follow it yourself)
+### Weighted microdata
+
+- Checked: `api.microdata.unweighted-series-are-labels`, `api.microdata.weight-stripping-confined`, `api.microdata.weights-not-aggregated-as-values`
+- `api.microdata.weights-not-values`: Microdata columns are weights, values, labels or counts: weights are only summed plainly to count entities or passed as weights=, values are aggregated with their own entity's weights, labels only group or filter, and counts derive from weights. (doc-only, not checked: follow it yourself)
+
### Code quality and types
- Checked: `api.quality.alembic-revisions-unused-imports`, `api.quality.ci-cd-scripts-ruff-exemptions`, `api.quality.command-line-modules-ruff-exemptions`, `api.quality.command-line-scripts-ruff-exemptions`, `api.quality.environment-variables-load`, `api.quality.mypy-v2`, `api.quality.ruff-format`, `api.quality.ruff-lint`, `api.quality.runtime-evaluated-models-resolve`, `api.quality.tests-pass`, `api.quality.tests-ruff-exemptions`, `api.quality.unittest-ruff-exemptions`
diff --git a/changelog.d/microdata-weight-rules.added.md b/changelog.d/microdata-weight-rules.added.md
new file mode 100644
index 000000000..12ad163f8
--- /dev/null
+++ b/changelog.d/microdata-weight-rules.added.md
@@ -0,0 +1 @@
+Added weighted-microdata rules to the registry, proven on the Simulation API's #767 household count: code in the package never aggregates a weight column as a value, never reads raw values from weighted microdata outside a designated module, and never reduces an unweighted label series. No current code uses microdata, so all three start enforced with no findings. A doc-only rule records that weights, values, labels and counts each allow different operations.
diff --git a/docs/engineering/skills/README.md b/docs/engineering/skills/README.md
index 709838f8a..d8b20e084 100644
--- a/docs/engineering/skills/README.md
+++ b/docs/engineering/skills/README.md
@@ -16,6 +16,8 @@ Current skills:
- `ci-workflows.md`: one line of code per workflow `run:` block, tested
workflow scripts, and shellcheck.
- `github-prs.md`: PR workflow and migration PR handoff expectations.
+- `microdata-weights.md`: weights, values, labels and counts in simulation
+ microdata, and the guards that keep weights from being aggregated as values.
- `migration_contracts.md`: API v2 migration route contracts, route-group
metadata, generated migration artifacts, and quality guards.
- `observability.md`: identifier ownership, HTTP and simulation transport,
diff --git a/docs/engineering/skills/microdata-weights.md b/docs/engineering/skills/microdata-weights.md
new file mode 100644
index 000000000..d52a41ba4
--- /dev/null
+++ b/docs/engineering/skills/microdata-weights.md
@@ -0,0 +1,82 @@
+# Microdata Weights Skill
+
+Read this before adding code that aggregates simulation microdata: anything
+that builds a `MicroSeries` or `MicroDataFrame`, or reads per-household or
+per-person columns from an economy result. No module in the package does
+today: the last one, the unused economy comparison, has been removed. The
+rules below keep new code honest.
+
+## Weights aren't values
+
+A microdata column is one of four kinds, and each allows different operations
+(rule `api.microdata.weights-not-values`):
+
+| Kind | Examples | Allowed | Never |
+| --- | --- | --- | --- |
+| Weight | `household_weight`, `person_weight` | A plain (unweighted) sum, to count entities; use as `weights=` | Aggregation as a value through a weighted series, frame or aggregator |
+| Value | `household_net_income`, `person_in_poverty` | Weighted aggregation, with its own entity's weights | Aggregation after the weights are stripped |
+| Label | `household_income_decile`, `age`, `is_male`, `race` | Grouping and filtering | Summing or averaging |
+| Count | `households`, people per decile | Derived from weights | |
+
+A weighted `MicroSeries` or `MicroDataFrame` weights every `.sum()`, `.mean()`
+and `.count()`. Aggregating a weight column that way weights it by itself and
+returns the sum of squared weights. In the Simulation API (#766, fixed by
+#767), that made `budget.households` about 10^4 times too high on every run,
+and its test missed it because it used plain `pd.DataFrame`s.
+
+policyengine's economy outputs arrive as `baseline` and `reform` dicts, whose
+microdata columns sit beside their weight columns. Code that builds a
+`MicroSeries` from them by hand follows three patterns:
+
+- Values pass `weights=`, and reform values reuse the baseline's weights
+ (`weights=baseline_income.weights`).
+- Labels (income and wealth deciles, `age`, `is_male`, `race`) are built
+ without `weights=` and are only grouping or filter keys.
+- Households are counted with a plain sum of `household_weight`.
+
+## Rules
+
+- `api.microdata.weights-not-aggregated-as-values`: a weight is never the value
+ of a weighted `MicroSeries`, the variable of a weighted aggregator, or
+ `[]` followed by a weighted reduction (`sum`, `mean`,
+ `median`, `quantile`, `std`, `count`, `groupby`, `cumsum`).
+- `api.microdata.weight-stripping-confined`: raw values are read from weighted
+ microdata (`np.array`, `np.asarray`, `.values`, `.to_numpy()`, `.tolist()`,
+ `pd.DataFrame(...)`, `pd.Series(...)`) only in a designated stripping module.
+ Summing stripped values gives an unweighted sample total, the mirror image of
+ #767. The Public API has no designated module yet.
+- `api.microdata.unweighted-series-are-labels`: a `MicroSeries` built without
+ `weights=` is never reduced, so it stays a grouping or filter key.
+
+The guards live in `tests/unit/services/test_microdata_weight_boundaries.py`
+and run in `make test`.
+
+## How the guard recognises a weight
+
+The guards read the source, so they recognise names, not types:
+
+- a weight is a string or f-string ending in `_weight` or `_weights` (including
+ `f"{entity}_weight"`), a name assigned from one, or a `.weights` attribute;
+- weighted microdata is a `MicroSeries` or `MicroDataFrame` built with
+ `weights=`, a name assigned from one, anything indexed, filtered or computed
+ from those, and, in a module that imports microdf, a string-keyed read from a
+ `baseline` or `reform` dict;
+- a label is a `MicroSeries` built without `weights=`, or a comparison or
+ filter of one.
+
+They don't follow data through function calls or containers, so a weight passed
+through a helper or a dict escapes them.
+
+## Handling a finding
+
+Each finding needs a decision, recorded in the PR:
+
+- **Fix it.** Count entities with a plain sum of the weights, aggregate values
+ through a weighted series, and reduce the weighted series grouped or filtered
+ by a label rather than the label itself.
+- **Move it.** A deliberate strip (for example, a per-household quantity used
+ only to sort households into buckets) moves into the designated stripping
+ module with a comment saying why, and the guard names that module.
+
+All three guards start with no findings and no baseline, so any finding
+fails the build until it's decided.
diff --git a/docs/generated/rules.md b/docs/generated/rules.md
index 227baa816..a2aed7157 100644
--- a/docs/generated/rules.md
+++ b/docs/generated/rules.md
@@ -8,16 +8,17 @@ Every rule that governs `policyengine-api` is listed here with how it is enforce
| Tier | enforced | ratcheted | advisory | doc-only | Total |
| --- | ---: | ---: | ---: | ---: | ---: |
-| mechanical | 48 | 0 | 1 | 12 | 61 |
+| mechanical | 51 | 0 | 1 | 13 | 65 |
| metric | 0 | 0 | 0 | 0 | 0 |
| judgment | 0 | 0 | 0 | 0 | 0 |
-**Blocking:** 48 of 61 rules.
+**Blocking:** 51 of 65 rules.
## Gaps
### Not enforced (doc-only)
+- `api.microdata.weights-not-values`: Microdata columns are weights, values, labels or counts: weights are only summed plainly to count entities or passed as weights=, values are aggregated with their own entity's weights, labels only group or filter, and counts derive from weights.
- `api.migration.alembic-autogenerated-only`: Alembic revisions are autogenerated against an explicitly selected v1 or v2 configuration; never hand-written DDL or data operations.
- `api.observability.id-is-diagnostic-only`: observability_id is used only to query diagnostic records, never for idempotency, cache ownership, authorization, database identity, routing or calculation behavior.
- `api.observability.registered-stages-only`: Span names come from the stage plans in policyengine_api.observability.stages; route and service code never adds span-name string literals.
@@ -264,6 +265,31 @@ flowchart LR
e_CI_job_v2_integration_check --> j_CI_pr_yml_v2_integration_check
```
+## Weighted microdata
+
+| Rule | Statement | Tier | Status | Blocking | Enforced by | CI | Proof |
+| --- | --- | --- | --- | --- | --- | --- | --- |
+| `api.microdata.unweighted-series-are-labels` | A MicroSeries built without weights= is a label: it only groups or filters weighted series and is never reduced itself. | mechanical | enforced | yes | pytest `tests/unit/services/test_microdata_weight_boundaries.py` | `pr.yml:test` | yes |
+| `api.microdata.weight-stripping-confined` | Raw values are read from weighted microdata (np.array, np.asarray, .values, .to_numpy(), .tolist(), pd.DataFrame or pd.Series) only in a designated stripping module, and the Public API has none yet. | mechanical | enforced | yes | pytest `tests/unit/services/test_microdata_weight_boundaries.py` | `pr.yml:test` | yes |
+| `api.microdata.weights-not-aggregated-as-values` | A weight column is never aggregated as a value: never the value of a weighted MicroSeries, the variable of a weighted aggregator, or frame[] followed by a weighted reduction. | mechanical | enforced | yes | pytest `tests/unit/services/test_microdata_weight_boundaries.py` | `pr.yml:test` | yes |
+| `api.microdata.weights-not-values` | Microdata columns are weights, values, labels or counts: weights are only summed plainly to count entities or passed as weights=, values are aggregated with their own entity's weights, labels only group or filter, and counts derive from weights. | mechanical | doc-only | no | nothing | — | — |
+
+```mermaid
+flowchart LR
+ r_api_microdata_unweighted_series_are_labels("api.microdata.unweighted-series-are-labels (enforced)")
+ e_pytest_tests_unit_services_test_microdata_weight_boundaries_py["pytest tests/unit/services/test_microdata_weight_boundaries.py"]
+ j_CI_pr_yml_test["CI pr.yml:test"]
+ r_api_microdata_weight_stripping_confined("api.microdata.weight-stripping-confined (enforced)")
+ r_api_microdata_weights_not_aggregated_as_values("api.microdata.weights-not-aggregated-as-values (enforced)")
+ r_api_microdata_weights_not_values("api.microdata.weights-not-values (doc-only)")
+ gap_unenforced>"not enforced"]
+ r_api_microdata_unweighted_series_are_labels --> e_pytest_tests_unit_services_test_microdata_weight_boundaries_py
+ e_pytest_tests_unit_services_test_microdata_weight_boundaries_py --> j_CI_pr_yml_test
+ r_api_microdata_weight_stripping_confined --> e_pytest_tests_unit_services_test_microdata_weight_boundaries_py
+ r_api_microdata_weights_not_aggregated_as_values --> e_pytest_tests_unit_services_test_microdata_weight_boundaries_py
+ r_api_microdata_weights_not_values -.-> gap_unenforced
+```
+
## Code quality and types
| Rule | Statement | Tier | Status | Blocking | Enforced by | CI | Proof |
diff --git a/rules/fixtures/api.microdata.unweighted-series-are-labels/bad_decile.py b/rules/fixtures/api.microdata.unweighted-series-are-labels/bad_decile.py
new file mode 100644
index 000000000..13b854936
--- /dev/null
+++ b/rules/fixtures/api.microdata.unweighted-series-are-labels/bad_decile.py
@@ -0,0 +1,20 @@
+"""Proof fixture: unweighted label series reduced. The rule must reject it."""
+
+from __future__ import annotations
+
+from microdf import MicroSeries
+
+ADULT_MIN_AGE = 18
+
+
+def average_decile(baseline: dict) -> float:
+ """Average the decile numbers themselves, as if they were incomes."""
+ decile = MicroSeries(baseline["household_income_decile"])
+ return float(decile.mean())
+
+
+def children(baseline: dict) -> float:
+ """Count the sample's child records, ignoring person weights."""
+ age = MicroSeries(baseline["age"])
+ is_child = age < ADULT_MIN_AGE
+ return float(is_child.sum())
diff --git a/rules/fixtures/api.microdata.unweighted-series-are-labels/good_decile.py b/rules/fixtures/api.microdata.unweighted-series-are-labels/good_decile.py
new file mode 100644
index 000000000..9544031ab
--- /dev/null
+++ b/rules/fixtures/api.microdata.unweighted-series-are-labels/good_decile.py
@@ -0,0 +1,28 @@
+"""Proof fixture: label series used only to group and filter.
+
+The rule must accept it.
+"""
+
+from __future__ import annotations
+
+from microdf import MicroSeries
+
+ADULT_MIN_AGE = 18
+
+
+def average_income_by_decile(baseline: dict) -> dict:
+ """Average household net income in each income decile, weighted by household."""
+ income = MicroSeries(
+ baseline["household_net_income"], weights=baseline["household_weight"]
+ )
+ decile = MicroSeries(baseline["household_income_decile"])
+ return income.groupby(decile).mean().to_dict()
+
+
+def child_poverty_rate(baseline: dict) -> float:
+ """Share of children in poverty, weighted by person."""
+ in_poverty = MicroSeries(
+ baseline["person_in_poverty"], weights=baseline["person_weight"]
+ )
+ age = MicroSeries(baseline["age"])
+ return float(in_poverty[age < ADULT_MIN_AGE].mean())
diff --git a/rules/fixtures/api.microdata.weight-stripping-confined/bad_net_income.py b/rules/fixtures/api.microdata.weight-stripping-confined/bad_net_income.py
new file mode 100644
index 000000000..14d7f18d2
--- /dev/null
+++ b/rules/fixtures/api.microdata.weight-stripping-confined/bad_net_income.py
@@ -0,0 +1,23 @@
+"""Proof fixture: weights stripped before a total. The rule must reject it.
+
+Summing the raw values of a weighted column gives an unweighted sample total,
+the mirror image of #767.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+from microdf import MicroSeries
+
+
+def baseline_net_income(baseline: dict) -> float:
+ """Total household net income over the sample's records, ignoring weights."""
+ return float(np.array(baseline["household_net_income"]).sum())
+
+
+def reform_net_income(reform: dict) -> float:
+ """Total household net income after dropping the weights it was given."""
+ income = MicroSeries(
+ reform["household_net_income"], weights=reform["household_weight"]
+ )
+ return float(income.values.sum())
diff --git a/rules/fixtures/api.microdata.weight-stripping-confined/good_net_income.py b/rules/fixtures/api.microdata.weight-stripping-confined/good_net_income.py
new file mode 100644
index 000000000..ae4f9bd16
--- /dev/null
+++ b/rules/fixtures/api.microdata.weight-stripping-confined/good_net_income.py
@@ -0,0 +1,24 @@
+"""Proof fixture: totals that keep their weights. The rule must accept it."""
+
+from __future__ import annotations
+
+from microdf import MicroSeries
+
+
+def baseline_net_income(baseline: dict) -> float:
+ """Total household net income, weighted by household."""
+ income = MicroSeries(
+ baseline["household_net_income"], weights=baseline["household_weight"]
+ )
+ return float(income.sum())
+
+
+def reform_net_income(baseline: dict, reform: dict) -> float:
+ """Total reform household net income, weighted by the baseline's households."""
+ baseline_income = MicroSeries(
+ baseline["household_net_income"], weights=baseline["household_weight"]
+ )
+ income = MicroSeries(
+ reform["household_net_income"], weights=baseline_income.weights
+ )
+ return float(income.sum())
diff --git a/rules/fixtures/api.microdata.weights-not-aggregated-as-values/bad_budget.py b/rules/fixtures/api.microdata.weights-not-aggregated-as-values/bad_budget.py
new file mode 100644
index 000000000..e7b6dbaa7
--- /dev/null
+++ b/rules/fixtures/api.microdata.weights-not-aggregated-as-values/bad_budget.py
@@ -0,0 +1,37 @@
+"""Proof fixture: a weight aggregated as a value. The rule must reject it.
+
+``build_budgetary_impact`` holds the Simulation API line that #767 fixed. The
+household frame is a MicroDataFrame weighted by ``household_weight``, so
+summing ``household_weight`` through it returns the sum of squared weights.
+``budgetary_impact`` makes the same mistake with a hand-built MicroSeries.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from microdf import MicroSeries
+from policyengine_simulation_executor.simulation_macro_output import BudgetaryImpact
+from policyengine_simulation_executor.simulation_output_common import (
+ _sum_output_variable,
+)
+
+if TYPE_CHECKING:
+ from policyengine.core import Simulation
+
+
+def build_budgetary_impact(baseline: Simulation) -> BudgetaryImpact:
+ """Count households as the Simulation API did before #767."""
+ return BudgetaryImpact(
+ households=_sum_output_variable(
+ baseline, "household_weight", entity="household"
+ ),
+ )
+
+
+def budgetary_impact(baseline: dict) -> dict:
+ """Count households by weighting the household weights by themselves."""
+ households = MicroSeries(
+ baseline["household_weight"], weights=baseline["household_weight"]
+ )
+ return {"households": households.sum()}
diff --git a/rules/fixtures/api.microdata.weights-not-aggregated-as-values/good_budget.py b/rules/fixtures/api.microdata.weights-not-aggregated-as-values/good_budget.py
new file mode 100644
index 000000000..1ee1a5a6a
--- /dev/null
+++ b/rules/fixtures/api.microdata.weights-not-aggregated-as-values/good_budget.py
@@ -0,0 +1,32 @@
+"""Proof fixture: entities counted with a plain sum of their weights.
+
+The rule must accept it. ``build_budgetary_impact`` holds #767's replacement
+line in the Simulation API, and ``budgetary_impact`` counts households with a
+plain sum of ``household_weight``.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from policyengine_simulation_executor.simulation_macro_output import BudgetaryImpact
+from policyengine_simulation_executor.simulation_output_common import (
+ _count_output_entities,
+)
+
+if TYPE_CHECKING:
+ from policyengine.core import Simulation
+
+
+def build_budgetary_impact(baseline: Simulation) -> BudgetaryImpact:
+ """Count households as #767 does: a plain sum of the household weights."""
+ return BudgetaryImpact(
+ households=_count_output_entities(
+ baseline, entity="household", weight="household_weight"
+ ),
+ )
+
+
+def budgetary_impact(baseline: dict) -> dict:
+ """Count households with a plain sum of their weights."""
+ return {"households": sum(baseline["household_weight"])}
diff --git a/rules/registry.toml b/rules/registry.toml
index 129099ca8..1f3059e0f 100644
--- a/rules/registry.toml
+++ b/rules/registry.toml
@@ -31,6 +31,10 @@ title = "Data source isolation"
id = "migration"
title = "Migration and schema"
+[[category]]
+id = "microdata"
+title = "Weighted microdata"
+
[[category]]
id = "quality"
title = "Code quality and types"
@@ -668,6 +672,83 @@ workflow = ".github/workflows/pr.yml"
job = "ensure-policyengine-bundle-supported-by-simulation-api"
invocation = "check-policyengine-bundle-supported.sh"
+# --- Weighted microdata ---------------------------------------------------
+
+[[rule]]
+id = "api.microdata.weights-not-values"
+category = "microdata"
+statement = "Microdata columns are weights, values, labels or counts: weights are only summed plainly to count entities or passed as weights=, values are aggregated with their own entity's weights, labels only group or filter, and counts derive from weights."
+rationale = "docs/engineering/skills/microdata-weights.md#weights-arent-values"
+scope = ["policyengine_api/**/*.py"]
+tier = "mechanical"
+status = "doc-only"
+autofix = false
+[rule.enforcer]
+kind = "none"
+
+[[rule]]
+id = "api.microdata.weights-not-aggregated-as-values"
+category = "microdata"
+statement = "A weight column is never aggregated as a value: never the value of a weighted MicroSeries, the variable of a weighted aggregator, or frame[] followed by a weighted reduction."
+rationale = "docs/engineering/skills/microdata-weights.md#rules"
+scope = ["policyengine_api/**/*.py"]
+tier = "mechanical"
+status = "enforced"
+autofix = false
+[rule.enforcer]
+kind = "pytest"
+nodes = ["tests/unit/services/test_microdata_weight_boundaries.py::test_weights_are_never_aggregated_as_values"]
+[[rule.ci]]
+workflow = ".github/workflows/pr.yml"
+job = "test"
+invocation = "make test"
+[rule.proof]
+bad_fixture = "rules/fixtures/api.microdata.weights-not-aggregated-as-values/bad_budget.py"
+good_fixture = "rules/fixtures/api.microdata.weights-not-aggregated-as-values/good_budget.py"
+test = "tests/unit/services/test_microdata_weight_boundaries.py::test_rule_proof_weights_not_aggregated_as_values"
+
+[[rule]]
+id = "api.microdata.weight-stripping-confined"
+category = "microdata"
+statement = "Raw values are read from weighted microdata (np.array, np.asarray, .values, .to_numpy(), .tolist(), pd.DataFrame or pd.Series) only in a designated stripping module, and the Public API has none yet."
+rationale = "docs/engineering/skills/microdata-weights.md#handling-a-finding"
+scope = ["policyengine_api/**/*.py"]
+tier = "mechanical"
+status = "enforced"
+autofix = false
+[rule.enforcer]
+kind = "pytest"
+nodes = ["tests/unit/services/test_microdata_weight_boundaries.py::test_weight_stripping_stays_in_the_designated_module"]
+[[rule.ci]]
+workflow = ".github/workflows/pr.yml"
+job = "test"
+invocation = "make test"
+[rule.proof]
+bad_fixture = "rules/fixtures/api.microdata.weight-stripping-confined/bad_net_income.py"
+good_fixture = "rules/fixtures/api.microdata.weight-stripping-confined/good_net_income.py"
+test = "tests/unit/services/test_microdata_weight_boundaries.py::test_rule_proof_weight_stripping_confined"
+
+[[rule]]
+id = "api.microdata.unweighted-series-are-labels"
+category = "microdata"
+statement = "A MicroSeries built without weights= is a label: it only groups or filters weighted series and is never reduced itself."
+rationale = "docs/engineering/skills/microdata-weights.md#rules"
+scope = ["policyengine_api/**/*.py"]
+tier = "mechanical"
+status = "enforced"
+autofix = false
+[rule.enforcer]
+kind = "pytest"
+nodes = ["tests/unit/services/test_microdata_weight_boundaries.py::test_unweighted_series_are_only_labels"]
+[[rule.ci]]
+workflow = ".github/workflows/pr.yml"
+job = "test"
+invocation = "make test"
+[rule.proof]
+bad_fixture = "rules/fixtures/api.microdata.unweighted-series-are-labels/bad_decile.py"
+good_fixture = "rules/fixtures/api.microdata.unweighted-series-are-labels/good_decile.py"
+test = "tests/unit/services/test_microdata_weight_boundaries.py::test_rule_proof_unweighted_series_are_labels"
+
# --- Code quality and types ------------------------------------------------
[[rule]]
diff --git a/tests/unit/services/test_microdata_weight_boundaries.py b/tests/unit/services/test_microdata_weight_boundaries.py
new file mode 100644
index 000000000..2fb72327b
--- /dev/null
+++ b/tests/unit/services/test_microdata_weight_boundaries.py
@@ -0,0 +1,511 @@
+"""Weighted-microdata guards: weights aren't values.
+
+Microdata columns come in four kinds (weights, values, labels and counts), and
+each allows different operations; docs/engineering/skills/microdata-weights.md
+sets them out. These checks read the source, so they recognise a weight by its
+name: a string or f-string ending in ``_weight`` or ``_weights``, a name
+assigned from one, or a ``.weights`` attribute. They don't follow data through
+function calls or containers.
+"""
+
+from __future__ import annotations
+
+import ast
+import enum
+from dataclasses import dataclass, field
+from functools import cache
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).parents[3]
+PACKAGE_ROOT = PROJECT_ROOT / "policyengine_api"
+FIXTURES_ROOT = PROJECT_ROOT / "rules" / "fixtures"
+
+WEIGHT_SUFFIXES = ("_weight", "_weights")
+# Reductions that a weighted MicroSeries or MicroDataFrame column weights.
+WEIGHTED_REDUCTIONS = frozenset(
+ {"sum", "mean", "median", "quantile", "std", "count", "groupby", "cumsum"}
+)
+# Weighted aggregators, with the position and keyword of the argument each one
+# aggregates as a value.
+WEIGHTED_AGGREGATORS: dict[str, tuple[int | None, str]] = {
+ "Aggregate": (None, "variable"),
+ "ChangeAggregate": (None, "variable"),
+ # The Simulation API's output helpers. The Public API has none, but the
+ # rule's proof is the real Simulation API line that #767 fixed.
+ "_sum_output_variable": (1, "variable"),
+ "_change_output_variable": (2, "variable"),
+}
+# Calls and attributes that return a weighted column's or frame's raw values.
+STRIPPING_FUNCTIONS = frozenset(
+ {
+ "np.array",
+ "np.asarray",
+ "numpy.array",
+ "numpy.asarray",
+ "pd.DataFrame",
+ "pd.Series",
+ "pandas.DataFrame",
+ "pandas.Series",
+ }
+)
+STRIPPING_METHODS = frozenset({"to_numpy", "tolist"})
+STRIPPING_ATTRIBUTES = frozenset({"values"})
+# Methods whose result is weighted exactly as their receiver is.
+WEIGHTING_PRESERVING_METHODS = frozenset(
+ {"abs", "astype", "clip", "copy", "fillna", "round"}
+)
+# policyengine's economy outputs arrive as ``baseline`` and ``reform`` dicts,
+# whose microdata columns sit beside their ``_weight`` columns. In a
+# module that imports microdf, a string-keyed read from either is a microdata
+# column.
+MICRODATA_DICT_NAMES = frozenset({"baseline", "reform"})
+# Modules allowed to read raw values from weighted microdata. The Public API has
+# none yet; a module that has to strip weights becomes the one designated here.
+WEIGHT_STRIPPING_MODULES: frozenset[str] = frozenset()
+
+AGGREGATION = "weight aggregated as a value"
+STRIPPING = "weights stripped"
+LABEL_REDUCTION = "unweighted series reduced"
+
+
+class Weighting(enum.Enum):
+ """What an expression holds, as far as weights go."""
+
+ WEIGHTED = "weighted"
+ LABEL = "label"
+
+
+@dataclass(frozen=True)
+class Finding:
+ """One place a check flags, with a line-independent baseline entry."""
+
+ check: str
+ path: str
+ line: int
+ function: str
+ source: str
+
+ @property
+ def entry(self) -> str:
+ """Return the baseline entry: path, enclosing function and expression."""
+ return f"{self.path}::{self.function}::{self.source}"
+
+ def __str__(self) -> str:
+ """Describe the finding with its line, for assertion messages."""
+ return f"{self.path}:{self.line} ({self.function}): {self.source}"
+
+
+@dataclass
+class Scope:
+ """The names a module or function assigns, sorted by what they hold."""
+
+ reads_microdata_dicts: bool
+ weight_keys: set[str] = field(default_factory=set)
+ weight_columns: set[str] = field(default_factory=set)
+ weight_arrays: set[str] = field(default_factory=set)
+ weighting: dict[str, Weighting] = field(default_factory=dict)
+
+ def child(self) -> Scope:
+ """Return a nested function's scope, which starts with these names."""
+ return Scope(
+ reads_microdata_dicts=self.reads_microdata_dicts,
+ weight_keys=set(self.weight_keys),
+ weight_columns=set(self.weight_columns),
+ weight_arrays=set(self.weight_arrays),
+ weighting=dict(self.weighting),
+ )
+
+
+def _dotted_name(node: ast.expr) -> str:
+ if isinstance(node, ast.Name):
+ return node.id
+ if isinstance(node, ast.Attribute):
+ return f"{_dotted_name(node.value)}.{node.attr}"
+ return ""
+
+
+def _called_name(call: ast.Call) -> str:
+ if isinstance(call.func, ast.Name):
+ return call.func.id
+ if isinstance(call.func, ast.Attribute):
+ return call.func.attr
+ return ""
+
+
+def _keyword(call: ast.Call, name: str) -> ast.expr | None:
+ for keyword in call.keywords:
+ if keyword.arg == name:
+ return keyword.value
+ return None
+
+
+def _argument(call: ast.Call, position: int | None, name: str) -> ast.expr | None:
+ if position is not None and len(call.args) > position:
+ return call.args[position]
+ return _keyword(call, name)
+
+
+def _names_a_weight(node: ast.expr) -> bool:
+ """Whether a string or f-string literal ends in ``_weight`` or ``_weights``."""
+ if isinstance(node, ast.Constant) and isinstance(node.value, str):
+ return node.value.endswith(WEIGHT_SUFFIXES)
+ if isinstance(node, ast.JoinedStr) and node.values:
+ tail = node.values[-1]
+ return (
+ isinstance(tail, ast.Constant)
+ and isinstance(tail.value, str)
+ and tail.value.endswith(WEIGHT_SUFFIXES)
+ )
+ return False
+
+
+def _is_weight_key(node: ast.expr, scope: Scope) -> bool:
+ if isinstance(node, ast.Name):
+ return node.id in scope.weight_keys
+ return _names_a_weight(node)
+
+
+def _is_weight_column(node: ast.expr, scope: Scope) -> bool:
+ """Whether the expression is ``[]``, or a name assigned one."""
+ if isinstance(node, ast.Name):
+ return node.id in scope.weight_columns
+ return isinstance(node, ast.Subscript) and _is_weight_key(node.slice, scope)
+
+
+def _is_weight_array(node: ast.expr, scope: Scope) -> bool:
+ if isinstance(node, ast.Name):
+ return node.id in scope.weight_arrays
+ return isinstance(node, ast.Attribute) and node.attr == "weights"
+
+
+def _is_weight(node: ast.expr, scope: Scope) -> bool:
+ return (
+ _is_weight_key(node, scope)
+ or _is_weight_column(node, scope)
+ or _is_weight_array(node, scope)
+ )
+
+
+def _combined(weightings: list[Weighting | None]) -> Weighting | None:
+ if Weighting.WEIGHTED in weightings:
+ return Weighting.WEIGHTED
+ if Weighting.LABEL in weightings:
+ return Weighting.LABEL
+ return None
+
+
+def _constructed_weighting(call: ast.Call) -> Weighting | None:
+ weights = _keyword(call, "weights")
+ has_weights = weights is not None and not (
+ isinstance(weights, ast.Constant) and weights.value is None
+ )
+ name = _called_name(call)
+ if name == "MicroSeries":
+ return Weighting.WEIGHTED if has_weights else Weighting.LABEL
+ if name == "MicroDataFrame" and has_weights:
+ return Weighting.WEIGHTED
+ return None
+
+
+def _call_weighting(call: ast.Call, scope: Scope) -> Weighting | None:
+ if _called_name(call) in {"MicroSeries", "MicroDataFrame"}:
+ return _constructed_weighting(call)
+ if (
+ isinstance(call.func, ast.Attribute)
+ and call.func.attr in WEIGHTING_PRESERVING_METHODS
+ ):
+ return _weighting(call.func.value, scope)
+ return None
+
+
+def _subscript_weighting(subscript: ast.Subscript, scope: Scope) -> Weighting | None:
+ if (
+ scope.reads_microdata_dicts
+ and isinstance(subscript.value, ast.Name)
+ and subscript.value.id in MICRODATA_DICT_NAMES
+ and isinstance(subscript.slice, ast.Constant)
+ and isinstance(subscript.slice.value, str)
+ ):
+ return Weighting.WEIGHTED
+ receiver = subscript.value
+ if isinstance(receiver, ast.Attribute) and receiver.attr in {"loc", "iloc"}:
+ receiver = receiver.value
+ return _weighting(receiver, scope)
+
+
+def _operands(node: ast.expr) -> list[ast.expr]:
+ """Return the operands whose weighting an operator's result keeps."""
+ if isinstance(node, ast.BinOp):
+ return [node.left, node.right]
+ if isinstance(node, ast.UnaryOp):
+ return [node.operand]
+ if isinstance(node, ast.Compare):
+ return [node.left, *node.comparators]
+ if isinstance(node, ast.BoolOp):
+ return list(node.values)
+ if isinstance(node, ast.IfExp):
+ return [node.body, node.orelse]
+ return []
+
+
+def _weighting(node: ast.expr, scope: Scope) -> Weighting | None:
+ """Return whether the expression holds weighted microdata or a label series."""
+ if isinstance(node, ast.Name):
+ return scope.weighting.get(node.id)
+ if isinstance(node, ast.Call):
+ return _call_weighting(node, scope)
+ if isinstance(node, ast.Subscript):
+ return _subscript_weighting(node, scope)
+ return _combined([_weighting(operand, scope) for operand in _operands(node)])
+
+
+def _assignments(body: list[ast.stmt]) -> list[tuple[str, ast.expr]]:
+ """Return each ``name = value`` in a body, outside nested functions and classes."""
+ assignments: list[tuple[str, ast.expr]] = []
+ pending: list[ast.AST] = list(body)
+ while pending:
+ node = pending.pop()
+ if isinstance(
+ node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)
+ ):
+ continue
+ if isinstance(node, ast.Assign):
+ assignments.extend(
+ (target.id, node.value)
+ for target in node.targets
+ if isinstance(target, ast.Name)
+ )
+ elif (
+ isinstance(node, (ast.AnnAssign, ast.NamedExpr))
+ and isinstance(node.target, ast.Name)
+ and node.value is not None
+ ):
+ assignments.append((node.target.id, node.value))
+ pending.extend(ast.iter_child_nodes(node))
+ return assignments
+
+
+def _with_assignments(scope: Scope, body: list[ast.stmt]) -> Scope:
+ """Sort the body's assigned names into the scope, until nothing changes."""
+ assignments = _assignments(body)
+ changed = True
+ while changed:
+ changed = False
+ for name, value in assignments:
+ groups = (
+ (scope.weight_keys, _is_weight_key(value, scope)),
+ (scope.weight_columns, _is_weight_column(value, scope)),
+ (scope.weight_arrays, _is_weight_array(value, scope)),
+ )
+ for names, matches in groups:
+ if matches and name not in names:
+ names.add(name)
+ changed = True
+ weighting = _combined([scope.weighting.get(name), _weighting(value, scope)])
+ if weighting is not None and scope.weighting.get(name) is not weighting:
+ scope.weighting[name] = weighting
+ changed = True
+ return scope
+
+
+class _MicrodataWalker:
+ """Walk a module with each function's scope and collect every finding."""
+
+ def __init__(self, relative_path: str) -> None:
+ self.relative_path = relative_path
+ self.findings: list[Finding] = []
+ self._called_functions: set[int] = set()
+
+ def walk(self, tree: ast.Module, *, reads_microdata_dicts: bool) -> None:
+ scope = _with_assignments(
+ Scope(reads_microdata_dicts=reads_microdata_dicts), tree.body
+ )
+ self._walk(tree, scope, "")
+
+ def _walk(self, node: ast.AST, scope: Scope, function: str) -> None:
+ for child in ast.iter_child_nodes(node):
+ if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
+ name = (
+ child.name if function == "" else f"{function}.{child.name}"
+ )
+ if isinstance(child, ast.ClassDef):
+ self._walk(child, scope, name)
+ else:
+ self._walk(
+ child, _with_assignments(scope.child(), child.body), name
+ )
+ continue
+ self._check(child, scope, function)
+ self._walk(child, scope, function)
+
+ def _report(self, check: str, node: ast.expr, function: str) -> None:
+ self.findings.append(
+ Finding(
+ check=check,
+ path=self.relative_path,
+ line=node.lineno,
+ function=function,
+ source=ast.unparse(node),
+ )
+ )
+
+ def _check(self, node: ast.AST, scope: Scope, function: str) -> None:
+ if isinstance(node, ast.Call):
+ self._called_functions.add(id(node.func))
+ self._check_aggregation(node, scope, function)
+ self._check_stripping_call(node, scope, function)
+ self._check_label_reduction(node, scope, function)
+ elif (
+ isinstance(node, ast.Attribute)
+ and node.attr in STRIPPING_ATTRIBUTES
+ and id(node) not in self._called_functions
+ and _weighting(node.value, scope) is Weighting.WEIGHTED
+ ):
+ self._report(STRIPPING, node, function)
+
+ def _check_aggregation(self, call: ast.Call, scope: Scope, function: str) -> None:
+ name = _called_name(call)
+ aggregated: ast.expr | None = None
+ if name in WEIGHTED_AGGREGATORS:
+ aggregated = _argument(call, *WEIGHTED_AGGREGATORS[name])
+ elif name == "MicroSeries" and _keyword(call, "weights") is not None:
+ aggregated = _argument(call, 0, "data")
+ if aggregated is not None and _is_weight(aggregated, scope):
+ self._report(AGGREGATION, call, function)
+ return
+ if (
+ isinstance(call.func, ast.Attribute)
+ and call.func.attr in WEIGHTED_REDUCTIONS
+ and _is_weight_column(call.func.value, scope)
+ ):
+ self._report(AGGREGATION, call, function)
+
+ def _check_stripping_call(
+ self, call: ast.Call, scope: Scope, function: str
+ ) -> None:
+ if _dotted_name(call.func) in STRIPPING_FUNCTIONS:
+ stripped = _argument(call, 0, "data")
+ elif (
+ isinstance(call.func, ast.Attribute) and call.func.attr in STRIPPING_METHODS
+ ):
+ stripped = call.func.value
+ else:
+ return
+ if stripped is not None and _weighting(stripped, scope) is Weighting.WEIGHTED:
+ self._report(STRIPPING, call, function)
+
+ def _check_label_reduction(
+ self, call: ast.Call, scope: Scope, function: str
+ ) -> None:
+ if (
+ isinstance(call.func, ast.Attribute)
+ and call.func.attr in WEIGHTED_REDUCTIONS
+ and _weighting(call.func.value, scope) is Weighting.LABEL
+ ):
+ self._report(LABEL_REDUCTION, call, function)
+
+
+def _imports_microdf(tree: ast.Module) -> bool:
+ for node in ast.walk(tree):
+ if isinstance(node, ast.ImportFrom) and (node.module or "").startswith(
+ "microdf"
+ ):
+ return True
+ if isinstance(node, ast.Import) and any(
+ alias.name.startswith("microdf") for alias in node.names
+ ):
+ return True
+ return False
+
+
+@cache
+def _module_findings(path: Path) -> tuple[Finding, ...]:
+ tree = ast.parse(path.read_text(encoding="utf-8"))
+ walker = _MicrodataWalker(path.relative_to(PROJECT_ROOT).as_posix())
+ walker.walk(tree, reads_microdata_dicts=_imports_microdf(tree))
+ return tuple(walker.findings)
+
+
+def _package_findings(check: str) -> list[Finding]:
+ modules = sorted(PACKAGE_ROOT.rglob("*.py"))
+ assert modules, f"no modules under {PACKAGE_ROOT}"
+ return [
+ finding
+ for module in modules
+ for finding in _module_findings(module)
+ if finding.check == check
+ ]
+
+
+def _fixture_sources(rule_id: str, fixture: str, check: str) -> list[str]:
+ path = FIXTURES_ROOT / rule_id / fixture
+ return [
+ finding.source for finding in _module_findings(path) if finding.check == check
+ ]
+
+
+def _described(findings: list[Finding]) -> str:
+ return "\n".join(str(finding) for finding in findings)
+
+
+def test_weights_are_never_aggregated_as_values() -> None:
+ findings = _package_findings(AGGREGATION)
+ assert findings == [], (
+ "a weight is aggregated as a value, which weights it by itself (the sum "
+ "of squared weights, as in #767); count entities with a plain sum of the "
+ "weights instead:\n" + _described(findings)
+ )
+
+
+def test_weight_stripping_stays_in_the_designated_module() -> None:
+ findings = [
+ finding
+ for finding in _package_findings(STRIPPING)
+ if finding.path not in WEIGHT_STRIPPING_MODULES
+ ]
+ assert findings == [], (
+ "raw values read from weighted microdata, outside the designated module; "
+ "an unweighted total of them is a sample total, the mirror image of #767. "
+ "Keep the weights, or decide with a reviewer to move the read into a "
+ "designated module:\n" + _described(findings)
+ )
+
+
+def test_unweighted_series_are_only_labels() -> None:
+ findings = _package_findings(LABEL_REDUCTION)
+ assert findings == [], (
+ "a MicroSeries built without weights= is reduced; such a series is a "
+ "label, only for grouping or filtering, so reduce the weighted series "
+ "grouped or filtered by it:\n" + _described(findings)
+ )
+
+
+def test_rule_proof_weights_not_aggregated_as_values() -> None:
+ """Proof for rule api.microdata.weights-not-aggregated-as-values."""
+ rule_id = "api.microdata.weights-not-aggregated-as-values"
+ assert _fixture_sources(rule_id, "bad_budget.py", AGGREGATION) == [
+ "_sum_output_variable(baseline, 'household_weight', entity='household')",
+ "MicroSeries(baseline['household_weight'], "
+ "weights=baseline['household_weight'])",
+ ]
+ assert _fixture_sources(rule_id, "good_budget.py", AGGREGATION) == []
+
+
+def test_rule_proof_weight_stripping_confined() -> None:
+ """Proof for rule api.microdata.weight-stripping-confined."""
+ rule_id = "api.microdata.weight-stripping-confined"
+ assert _fixture_sources(rule_id, "bad_net_income.py", STRIPPING) == [
+ "np.array(baseline['household_net_income'])",
+ "income.values",
+ ]
+ assert _fixture_sources(rule_id, "good_net_income.py", STRIPPING) == []
+
+
+def test_rule_proof_unweighted_series_are_labels() -> None:
+ """Proof for rule api.microdata.unweighted-series-are-labels."""
+ rule_id = "api.microdata.unweighted-series-are-labels"
+ assert _fixture_sources(rule_id, "bad_decile.py", LABEL_REDUCTION) == [
+ "decile.mean()",
+ "is_child.sum()",
+ ]
+ assert _fixture_sources(rule_id, "good_decile.py", LABEL_REDUCTION) == []