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) == []