From 0835944512f8632f878168bbe1d2eceba951500d Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 5 Nov 2026 04:39:29 +0000 Subject: [PATCH 01/11] Guard against random draws in the engine Port policyengine-us's scanner for engine-time randomness and hash draws as a code-health test. EXEMPT lists the package's SHA-256 digests, which fingerprint, key caches or check content; RATCHET lists today's draws: the splitmix64 helper, the formulas and uprating that call it, and the seeded numpy draws in participation and two scenarios. Cases that need newer syntax are filtered at collection, so no Python reports a skip. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LwQq6H1FBndAKeuXdSGuft --- .../test_no_engine_time_randomness.py | 4203 +++++++++++++++++ 1 file changed, 4203 insertions(+) create mode 100644 policyengine_uk/tests/code_health/test_no_engine_time_randomness.py diff --git a/policyengine_uk/tests/code_health/test_no_engine_time_randomness.py b/policyengine_uk/tests/code_health/test_no_engine_time_randomness.py new file mode 100644 index 0000000000..5b38e94666 --- /dev/null +++ b/policyengine_uk/tests/code_health/test_no_engine_time_randomness.py @@ -0,0 +1,4203 @@ +"""The model draws nothing while a simulation is built or run. + +A formula is a pure, deterministic function of its inputs +(policyengine-core#500), and seeding does not exempt a draw. Stochastic +inputs such as take-up flags are drawn when the dataset is built and stored +as input variables; the model compares them with its rules. policyengine-core#519 checks formulas +for numpy and stdlib randomness when variables register, but a hand-written +integer hash of entity ids passes that check, so this scan looks for the +hash as well. + +Every module of the package except tests is parsed (formulas, reforms, +scenarios, simulation code, data loading, uprating and utilities). A name resolves to +everything it may refer to where it is used: the targets of every binding +that may be the one read, in each scope Python looks in, so the scan errs +toward reporting. Bindings come from imports, named imports and re-exports +followed across the package, star imports, assignments of a name or +attribute chain (``xp = np``, also in tuples), default arguments, ``global`` +and ``nonlocal`` (the nearest enclosing function that binds the name, +wherever in its body; either sends the name there in the code nested in +its scope too, unless that code binds the name itself), functions, lambdas +and comprehensions. An attribute +of a package module, or of a module EXTERNAL_EXPORTS maps, is what that +module binds under the name, so ``api.np.random`` is numpy's random module +where ``api`` is ``policyengine_core.model_api``. No code sees the names a +class body it is nested in binds, or its ``global`` and ``nonlocal`` +statements. A module or class body looks further out +when no binding is sure to have run before the use with no ``del`` or +``except ... as`` cleanup that may come between (all are in an ``if``, loop, +``try``, ``with``, ``match`` or ``assert``, in a comprehension, past the +first operand of ``and``, ``or`` or a chained comparison, or in a branch of +a conditional expression, or the use comes first). A class body that binds +the name anywhere (an annotation alone or a ``del`` counts) then looks in +the module, never in a function around it, and so does what the body +evaluates: its annotations, its methods' defaults and annotations, and +type-parameter bounds and ``type`` values in it. A function may run +as soon as its ``def`` has run, so the scopes around it count only the +bindings before the ``def``, and every later deletion. A function's ``del`` +of a name it declares ``global`` runs whenever the function is called, so +after the ``def`` it may come between any binding and any use. A generator +expression runs its body as it is iterated, so the same holds for it from +the expression (its first iterable runs at once). A class body, +a decorated function and a lambda may run before the statement that creates +them has bound its names (a decorator may call the function, and an +assignment's value may call a lambda in it), so code inside them doesn't +take those bindings as sure. Nor does a use take as sure a binding its own +statement makes (a multi-line assignment, or a default or base on a later +line of its ``def`` or ``class``): Python makes those after evaluating it, +and treating a walrus earlier in the statement the same way errs toward +reporting. An annotation is read as code that may run any time after its +statement starts, seeing a class body around it, which covers both ways +Python may evaluate it: with its statement, as up to 3.13, and whenever the +annotations are read, as 3.14 defers them (PEP 649). A type parameter's +bound or default and a ``type`` alias's value, which Python evaluates when +they are read (PEP 695, 3.12+; defaults PEP 696, 3.13+), are read the same +way. A ``type`` statement +binds its alias's name, but a type parameter's own name is not taken as +bound, which errs toward reporting. An annotation alone binds no value in a +module or class body. + +So ``from numpy import random as r`` followed by ``r.random(3)`` is caught, +as are ``xp = np`` followed by ``xp.random.uniform()``, ``from +policyengine_uk.model_api import np as lib`` followed by +``lib.random.uniform()``, and a ``try: import numpy as lib`` with a +fallback, while a local variable called ``random`` is not. A star import of +a package module binds its ``__all__`` (built from string literals with +``=``, ``+``, ``+=``, ``append``, ``extend``, ``insert`` or ``remove``; any +other ``__all__`` exports every name) or, without one, every name it binds +without a leading underscore. It may not bind a name an earlier or removed +``__all__`` held (an import cycle can see that one), one the module binds +only conditionally or deletes, or any name where ``__all__`` is built +otherwise, so the scan also looks further out. Where the module, while it +runs, may import the importing one (directly or through other modules, +counting every import in them, even in a function; the submodules a star +import of a package loads: the module each name its ``__all__`` lists +names below the package (``sub.leaf`` too, with ``sub``) or, where +``__all__`` is built otherwise, every module below it, and the same where +a literal ``__import__`` ``fromlist`` holds ``*`` (read as Python iterates +it: a list's, tuple's or set's elements, a dict's keys or a string's +characters, a literal ``*`` or ``**`` entry's own in its place); any call +whose literal arguments name a module as ``importlib.import_module`` or +``__import__`` would read them, whatever name the call is made under, a +relative one against its literal ``package`` (or the module's +``__package__`` or ``__name__``) or the package its ``globals`` give, as +Python reads them (``globals()`` or a dict literal, a literal ``**`` +entry's own entries in its place and the last of a repeated key counting: +``__package__``, else ``__spec__``, else ``__name__`` and ``__path__``); +the packages an import starts; the code of modules it may +reach without importing them, as another module may have loaded them: an +attribute chain, ``getattr`` with a literal name, ``sys.modules`` with a +literal key, or the module a literal dynamic import returns, from a name +any import, or assignment (or walrus) of such a chain or a conditional +choice of one (``if``-``else``, ``and``, ``or``), in the module binds, a +submodule a star import of a package may bind, or in a package's +``__init__`` its own submodule; and the imports in code that may run +later, when it is called or read (functions, lambdas, generator +expressions, annotations, type parameters and ``type`` values), in the +packages around a module, which have run before it, and in the modules +they import or reach as they run, since its code may call that code), the +star import may see it part-run, before it binds a name or sets ``__all__``: it may then bind any +name the module binds without a leading underscore, and binds none surely. +EXTERNAL_EXPORTS lists the +names numpy and ``policyengine_core.model_api`` bind where they matter here; +``np``, ``numpy`` and ``pd`` are bound to numpy and pandas wherever nothing +else binds them. Reported, as (file, enclosing function or class, kind): +- ``stochastic import``: importing ``random``, ``secrets`` or + ``numpy.random`` (or anything from them), including through + ``importlib.import_module`` or ``__import__`` with literal arguments, + positional or keyword (``package``, ``fromlist`` and ``level`` too); +- ``random source``: a name or attribute that resolves into ``numpy.random``, + ``random`` or ``secrets``, or to ``uuid.uuid1``, ``uuid.uuid4`` or + ``os.urandom``, including through ``getattr`` with a literal name; +- ``builtin hash``: any use of the unshadowed ``hash``, which Python salts + per process; +- ``hash library``: importing a hashing or checksum library (``hashlib``, + ``zlib``, ``uuid``, ...); +- ``hash use``: a name or attribute that resolves into one of those + libraries or to pandas' hashing functions, so that a new use of an + import EXEMPT allows still counts; +- ``hash constant``: an integer constant of a common hash mixer + (splitmix64, murmur3, FNV, xxHash, PCG), which catches a copied hash; +- ``xorshift``: the step ``x ^ (x >> k)`` such mixers are built from, + written with operators, ``^=`` or numpy's ``bitwise_xor`` and shifts; +- ``sampling call``: ``.sample(...)`` with sampling arguments, or a method + named ``shuffle``, ``permutation``, ``choice`` and the like; +- ``id modulo``: ``... % n``, ``%=``, ``np.mod``, ``divmod`` and the like, + where the dividend reads an id (a name or string ending ``_id``); +- ``draw helper call``: a use, under any name it is imported or re-exported + as, of a top-level function of the package (one defined in an ``if`` or + ``try`` too) whose body has a finding EXEMPT does not cover, or that + itself uses such a function, so that calling an existing draw again also + counts. + +This is a ratchet over recognised patterns, not a proof that no +hand-written draw exists. It does not follow ``eval`` or ``exec``, imports +of a computed name (a relative dynamic import whose ``package`` or +``globals`` is computed is read against the module's own package), +package search paths changed as the code runs (a ``__path__`` assignment) +or import hooks, code outside the package (what it imports, or calls back +into), values reached through objects or containers (``self.r = +np.random`` is caught where ``np.random`` is read; a later +``self.r.uniform()`` is not; nor are functions one module sets on another +or registers for a later call), names bound otherwise than by an import or +an assignment to a plain name (unpacking, ``for``, ``with``, parameters), +the values calls return (``r = +importlib.import_module("random")`` is caught at the import; a later +``r.random()`` is not), names ``match`` patterns capture, or methods as draw +helpers. A lookup that nests deeper than MAX_RESOLUTION_DEPTH fails the scan +rather than passing unresolved. + +EXEMPT lists reviewed occurrences that are not draws; RATCHET lists today's +engine-time draws that still have to move into the data. Both count exact +occurrences per (file, enclosing scope, kind): an extra occurrence fails, and +so does an entry that no longer matches as many as it says, so entries are +lowered and removed as sites go. RATCHET_CEILING only ever falls, and +RATCHET must reach zero. +""" + +import ast +from collections import Counter +from functools import cache +from itertools import permutations +from pathlib import Path +import sys +from typing import NamedTuple + +import pytest + + +PACKAGE = Path(__file__).resolve().parents[2] +TESTS = PACKAGE / "tests" + +RATCHET = { + # The splitmix64 finalizer that maps entity ids to uniform draws on + # [0, 1): its three mixing constants and three xorshift steps. Every + # other entry calls it, or calls what calls it. + "utils/stochastic.py": { + ("splitmix64_uniform", "hash constant"): 3, + ("splitmix64_uniform", "xorshift"): 3, + }, + # Formula: hash of benunit_id (salt 0) that decides whether a UC + # deduction applies and how large it is, in simulations built from data. + "variables/gov/dwp/universal_credit/deductions/uc_deduction_random_draw.py": { + ("uc_deduction_random_draw.formula", "draw helper call"): 1, + }, + # Formula: hash of benunit_id (salt 1) that picks the UC deduction type + # combination, in simulations built from data. + "variables/gov/dwp/universal_credit/deductions/uc_deduction_type_random_draw.py": { + ("uc_deduction_type_random_draw.formula", "draw helper call"): 1, + }, + # Formula: hash of household_id (salt 2) that decides whether a household + # on a qualifying benefit receives the Warm Home Discount, in simulations + # built from data. + "variables/gov/ofgem/warm_home_discount_random_draw.py": { + ("warm_home_discount_random_draw.formula", "draw helper call"): 1, + }, + # Formula: hash of person_id (salt 2) that orders people within each year + # of age and sex before stratified_uniform spreads them over the year, in + # simulations built from data. + "variables/household/demographic/months_since_last_birthday.py": { + ("months_since_last_birthday.formula", "draw helper call"): 1, + }, + # Uprating: uprate_student_loan_plans hashes person_id (salt 3) and + # assigns a plan to new borrowers whose draw falls below the take-up + # rate. apply_single_year_uprating calls it, apply_uprating calls that, + # and extend_single_year_dataset calls apply_uprating. + "data/economic_assumptions.py": { + ("uprate_student_loan_plans", "draw helper call"): 1, + ("apply_single_year_uprating", "draw helper call"): 1, + ("apply_uprating", "draw helper call"): 1, + ("extend_single_year_dataset", "draw helper call"): 1, + }, + # Simulation loading: extends a single-year dataset through + # extend_single_year_dataset, and so the student loan draw. + "simulation.py": { + ("Simulation._extend_dataset_with_cgt_provenance", "draw helper call"): 1, + }, + # Labour supply: a seeded numpy RandomState that draws participation + # responses. + "dynamics/participation.py": { + ("apply_participation_responses", "random source"): 1, + }, + # UC rebalancing scenario: a seeded numpy Generator draws one value per + # benefit unit and compares it with the new claimant share; the Scenario + # at module level passes on the modifier. + "scenarios/uc_reform.py": { + ("add_universal_credit_reform", "random source"): 1, + ("", "draw helper call"): 1, + }, + # PIP phase-in scenario: seeds numpy's global stream and draws one value + # per person, which removes PIP where it falls below the phase-in share; + # the Scenario at module level passes on the modifier. + "scenarios/pip_reform.py": { + ("modify_simulation", "random source"): 2, + ("", "draw helper call"): 1, + }, +} +# The number of RATCHET occurrences. Lower it when one goes; never raise it. +RATCHET_CEILING = 21 + +EXEMPT = { + # SHA-256 fingerprint of the files that shape the data build, for + # provenance. + "build_metadata.py": { + ("", "hash library"): 1, + ("get_data_build_fingerprint", "hash use"): 1, + }, + # SHA-256 of the capital gains projection's metadata and provenance, for + # identifying the projection. + "data/capital_gains_projection.py": { + ("", "hash library"): 1, + ("_fingerprint", "hash use"): 1, + }, + # SHA-256 of a dataset's bucket, path and generation, naming its local + # cache directory. + "data/dataset_sources.py": { + ("", "hash library"): 1, + ("_cached_dataset_path", "hash use"): 1, + }, + # SHA-256 of the uprating parameters, which the simulation compares for + # equality to decide whether a cached projection still applies. + "data/economic_assumptions.py": { + ("", "hash library"): 1, + ("uprating_input_fingerprint", "hash use"): 1, + }, + # Aggregate impact tool: SHA-256 of a file it records. + "utils/aggregate_impact/__main__.py": { + ("", "hash library"): 1, + ("_sha256", "hash use"): 1, + }, + # Aggregate impact tool: SHA-256 content check of a pinned dataset file. + "utils/aggregate_impact/dataset.py": { + ("", "hash library"): 1, + ("sha256_file", "hash use"): 1, + }, + # Aggregate impact tool: SHA-256 identities of the tool's files, the + # measured spec and reforms, the model's files and the installed + # distributions, stamped on a result. + "utils/aggregate_impact/measure.py": { + ("", "hash library"): 1, + ("_digest_tool_files", "hash use"): 1, + ("definitions_digest", "hash use"): 1, + ("model_fingerprint", "hash use"): 1, + ("runtime_identity", "hash use"): 1, + }, + # Aggregate impact tool: SHA-256 of a partition file, recorded with it. + "utils/aggregate_impact/partition.py": { + ("", "hash library"): 1, + ("load_partition", "hash use"): 1, + }, + # Offline LHA rates builder: SHA-256 content check of a downloaded source + # file. + "utils/build_lha_published_rates.py": { + ("", "hash library"): 1, + ("fetch", "hash use"): 1, + }, +} +OFFLINE_MODULES = set() + +# Names modules outside the package bind, where they matter here, for star +# and named imports alike. numpy's __all__ includes its random module. +EXTERNAL_EXPORTS = { + "numpy": {"random": "numpy.random"}, + "policyengine_core.model_api": { + "np": "numpy", + "numpy": "numpy", + "pd": "pandas", + # Core's guard that raises on formula-time randomness. + "random": "policyengine_core.commons.formulas.random", + }, +} +# What a name bound nowhere in the module refers to. +UNBOUND_NAMES = { + "np": "numpy", + "numpy": "numpy", + "pd": "pandas", + "hash": "builtins.hash", + "getattr": "builtins.getattr", + "__import__": "builtins.__import__", + "divmod": "builtins.divmod", +} +DYNAMIC_IMPORTS = {"importlib.import_module", "builtins.__import__"} +XOR_FUNCTIONS = {"numpy.bitwise_xor", "operator.xor"} +SHIFT_FUNCTIONS = { + "numpy.right_shift", + "numpy.left_shift", + "operator.rshift", + "operator.lshift", +} +MOD_FUNCTIONS = { + "numpy.mod", + "numpy.remainder", + "numpy.fmod", + "operator.mod", + "builtins.divmod", +} +STOCHASTIC_MODULES = ("numpy.random", "random", "secrets") +STOCHASTIC_CALLABLES = {"uuid.uuid1", "uuid.uuid4", "os.urandom"} +HASH_MODULES = { + "hashlib", + "zlib", + "binascii", + "uuid", + "mmh3", + "xxhash", + "blake3", + "crcmod", +} +HASH_APIS = ( + "pandas.util.hash_pandas_object", + "pandas.util.hash_array", + "pandas.core.util.hashing", +) +SAMPLING_METHODS = { + "shuffle", + "permutation", + "permuted", + "choice", + "choices", + "randint", + "randrange", +} +SAMPLING_KEYWORDS = {"n", "frac", "random_state", "weights", "replace"} +HASH_CONSTANTS = { + # splitmix64: golden gamma and the two finalizer multipliers + 0x9E3779B97F4A7C15, + 0xBF58476D1CE4E5B9, + 0x94D049BB133111EB, + # murmur3 fmix64 and fmix32 + 0xFF51AFD7ED558CCD, + 0xC4CEB9FE1A85EC53, + 0x85EBCA6B, + 0xC2B2AE35, + # FNV-1a 64-bit and 32-bit offset bases and primes + 0xCBF29CE484222325, + 0x100000001B3, + 0x811C9DC5, + 0x01000193, + # xxHash64 primes + 0x9E3779B185EBCA87, + 0xC2B2AE3D27D4EB4F, + 0x165667B19E3779F9, + 0x85EBCA77C2B2AE63, + 0x27D4EB2F165667C5, + # PCG and Knuth MMIX 64-bit LCG multiplier + 0x5851F42D4C957F2D, +} +FUNCTION_NODES = (ast.FunctionDef, ast.AsyncFunctionDef) +COMPREHENSION_NODES = (ast.ListComp, ast.SetComp, ast.GeneratorExp, ast.DictComp) +TYPE_ALIAS = getattr(ast, "TypeAlias", ()) # Python 3.12+ +# Statements whose bodies may not run, so a binding inside may not happen. +COMPOUND_NODES = ( + ast.If, + ast.For, + ast.AsyncFor, + ast.While, + ast.Try, + ast.TryStar, + ast.With, + ast.AsyncWith, + ast.Match, +) +# Lookups nested deeper than this are cut, and the package scan fails +# (test_package_lookups_stay_within_the_depth_limit). +MAX_RESOLUTION_DEPTH = 64 +# What ``from module import *`` binds when the module has no __all__, and +# when its __all__ is built in a way this scan does not evaluate. +DEFAULT_PUBLIC = "names without a leading underscore" +EVERY_NAME = "every name" + + +def _matches(target: str, prefixes) -> bool: + return any( + target == prefix or target.startswith(prefix + ".") for prefix in prefixes + ) + + +def _is_stochastic(target: str | None) -> bool: + if not target: + return False + return target in STOCHASTIC_CALLABLES or _matches(target, STOCHASTIC_MODULES) + + +def _classify(target: str) -> str | None: + """The kind of finding a use of ``target`` is, if any.""" + if _is_stochastic(target): + return "random source" + if target == "builtins.hash": + return "builtin hash" + if target.split(".")[0] in HASH_MODULES or _matches(target, HASH_APIS): + return "hash use" + return None + + +def _absolute(module: str, level: int, package: str | None) -> str: + """The absolute name of a ``from`` import's module.""" + if not level: + return module + if package is None: + return "." * level + module + parts = package.split(".") + parent = ".".join(parts[: len(parts) - level + 1]) + return f"{parent}.{module}" if module else parent + + +def _dotted(node): + """('np', ['random', 'seed']) for np.random.seed, else None.""" + parts = [] + while isinstance(node, ast.Attribute): + parts.append(node.attr) + node = node.value + if isinstance(node, ast.Name): + return node.id, list(reversed(parts)) + return None + + +def _literal(node) -> str | None: + if isinstance(node, ast.Constant) and isinstance(node.value, str): + return node.value + return None + + +def _targets(target) -> set: + """The dotted names a binding's target may be.""" + if target is None: + return set() + if isinstance(target, str): + return {target} + return set(target.targets()) + + +class _Alias: + """A name bound, at ``line``, to another name or attribute chain.""" + + def __init__(self, scope, value, line): + self.scope = scope + self.root, self.attributes = _dotted(value) + self.line = line + + def targets(self): + # The value is evaluated where and when the binding is made. + roots = self.scope.modules.guarded( + self, + frozenset(), + lambda: frozenset(self.scope.resolve(self.root, self.line)), + ) + return self.scope.modules.expand_all( + ".".join([root] + self.attributes) for root in roots + ) + + +class _Scope: + """One namespace: a module, class, function, lambda or comprehension, or + the annotations of one statement.""" + + def __init__(self, kind, qualname, parent, modules, line=None): + self.kind = kind + self.qualname = qualname + self.parent = parent + self.modules = modules + # Where the scope starts: a function's def, for instance. + self.line = line + # name -> [(line, target, conditional)]. The target is what the name + # refers to (a module, function or other import), an _Alias to + # resolve, or None for a value. A conditional binding + # (in an if, loop, try, with or match, or a deletion) may not happen. + self.bindings = {} + # name -> [line] of each del or except ... as cleanup that may unbind + # it. + self.deletions = {} + # name -> [line] of each def whose function may unbind it when + # called, which may be at any time after that line. + self.deletions_after = {} + # For a scope that may run before the statement creating it has bound + # its names (a class body, a decorated function, a lambda), that + # statement's (first line, last line). + self.unfinished = None + # (line, module, conditional, importing module) for each star + # import. + self.stars = [] + # Names a global or nonlocal statement gives to another scope. + self.declared = {} + # For a function, lambda or class body, the names Python makes + # local to it. + self.local_names = frozenset() + + def bind(self, name, line, target=None, conditional=False): + owner = self.declared.get(name) + if owner is not None: + # Bound in that scope whenever this one runs, which may be never. + return owner.bind(name, line, target, conditional=True) + self.bindings.setdefault(name, []).append((line, target, conditional)) + + def delete(self, name, line): + """Record that ``name`` may be unbound at ``line``.""" + self.bind(name, line, None, conditional=True) + owner = self.declared.get(name, self) + # A deletion in a function, of a name another scope binds, happens + # when the function is called: in the owner's time, any time after + # the def of the outermost function between them. + called_after = None + scope = self + while scope is not owner: + if scope.kind in ("function", "lambda"): + called_after = scope.line + scope = scope.parent + if called_after is None: + owner.deletions.setdefault(name, []).append(line) + else: + owner.deletions_after.setdefault(name, []).append(called_after) + + def own(self, name, line, deferred, unfinished=()): + """(targets, bound, falls through) for ``name`` in this scope alone. + + The use is at ``line`` or, with ``deferred``, at any time after it; + a None line is once the scope has run. A binding within one of the + ``unfinished`` (first line, last line) spans may not have happened + yet. + """ + entries = [(entry[0], entry[2]) for entry in self.bindings.get(name, [])] + targets = set() + for _, target, _ in self.bindings.get(name, []): + targets |= _targets(target) + for star_line, base, conditional, importer in self.stars: + provided = self.modules.star_binds(base, name, importer) + if provided is not None: + targets |= provided + sure = self.modules.star_binds_surely(base, name, importer) + entries.append((star_line, conditional or not sure)) + if not entries: + return targets, False, True + if self.kind not in ("module", "class"): + # A name bound anywhere in a function is local to all of it. + return targets, True, False + # A module or class body looks further out unless a binding is sure + # to have run before the use, with no deletion that may come between. + deletions = self.deletions.get(name, []) + deletions_after = self.deletions_after.get(name, []) + + def holds(bound_line): + if deferred or line is None: + return ( + (line is None or bound_line <= line) + and all(deleted < bound_line for deleted in deletions) + and not deletions_after + ) + return ( + bound_line < line + and not any(bound_line <= deleted <= line for deleted in deletions) + and not any(after <= line for after in deletions_after) + ) + + sure = [ + entry_line + for entry_line, conditional in entries + if not conditional + and not any(first <= entry_line <= last for first, last in unfinished) + ] + return targets, True, not any(map(holds, sure)) + + def resolve(self, name, line=None, deferred=False, statement=None) -> set: + """Every dotted target ``name`` may refer to at ``line``. + + The answer errs toward reporting: the targets of every binding that + may be the one read, in each scope Python looks in until a binding is + sure to hold. ``deferred`` says the code may run at any time after + ``line``. ``statement`` is the (first line, last line) of the + statement the use is in, whose own bindings may not have happened. + """ + targets = set() + # The scope a global or nonlocal statement sends the name to, until + # the lookup reaches it. + owner = None + first = True + unfinished = [statement] if statement else [] + scope = self + while scope is not None: + if scope.kind == "annotation": + # Up to 3.13 an annotation is evaluated with its statement. + # 3.14 defers it to whenever the annotations are read (PEP + # 649), any time after the statement starts, in a scope that + # also sees a class body around it. Reading it as deferred + # from the statement covers both, since no binding the + # statement makes is sure (resolve). + line, deferred = scope.line, True + scope = scope.parent + continue + if scope is owner: + owner = None + # A global or nonlocal name skips the scopes up to the one it + # names, and no code nested in a class body sees the names it + # binds or declares. + if owner is None and (first or scope.kind != "class"): + # A global or nonlocal statement here sends the name there, + # here and in the code nested here. + owner = scope.declared.get(name) + if owner is None: + found, bound, falls = scope.own(name, line, deferred, unfinished) + targets |= found + if bound and not falls: + return targets + if scope.kind == "class" and name in scope.local_names: + # A class body that binds the name reads it from its + # namespace, then the module's (LOAD_NAME): the + # functions around it are skipped. + owner = scope + while owner.parent is not None: + owner = owner.parent + if scope.unfinished is not None: + # It may run before its statement's bindings have happened. + unfinished.append(scope.unfinished) + if scope.kind in ("function", "lambda", "generator"): + # A function may run as soon as its def has run, and a + # generator expression's body as soon as it is made, so the + # scopes around them are read as of then, and any later + # deletion counts. + line, deferred = scope.line, True + first = False + scope = scope.parent + if name in UNBOUND_NAMES: + targets.add(UNBOUND_NAMES[name]) + return targets + + +def _comprehension_walrus_targets(node) -> set: + """Names a walrus binds anywhere in a comprehension, which Python binds + in the function around it. A lambda inside is a scope of its own, apart + from its defaults.""" + targets = set() + stack = [node] + while stack: + child = stack.pop() + if isinstance(child, ast.NamedExpr): + targets.add(child.target.id) + if isinstance(child, ast.Lambda): + stack.extend(child.args.defaults) + stack.extend(d for d in child.args.kw_defaults if d) + continue + stack.extend(ast.iter_child_nodes(child)) + return targets + + +def _function_locals(node) -> frozenset: + """The names Python makes local to a function or lambda, wherever its + body binds them: its parameters and every name bound outside nested + scopes, less those a global or nonlocal statement sends elsewhere.""" + arguments = node.args + parameters = arguments.posonlyargs + arguments.args + arguments.kwonlyargs + parameters += [arguments.vararg, arguments.kwarg] + names = {argument.arg for argument in parameters if argument is not None} + body = node.body if isinstance(node.body, list) else [node.body] + return frozenset(names) | _block_locals(body) + + +def _block_locals(body) -> frozenset: + """The names Python makes local to a block of statements, wherever it + binds them: every name bound outside nested scopes, an annotation alone + and a ``del`` included, less those a global or nonlocal statement sends + elsewhere.""" + names, declared = set(), set() + stack = list(body) + while stack: + child = stack.pop() + if isinstance(child, (ast.Global, ast.Nonlocal)): + declared.update(child.names) + continue + if isinstance(child, FUNCTION_NODES + (ast.ClassDef, ast.Lambda)): + # Only its name, decorators, defaults, annotations and bases are + # this scope's. Up to 3.13 annotations are evaluated here, so a + # walrus in one binds here; 3.14 defers them and refuses a walrus + # in one, so the earlier rule covers both. + if not isinstance(child, ast.Lambda): + names.add(child.name) + stack.extend(getattr(child, "decorator_list", [])) + if isinstance(child, ast.ClassDef): + stack.extend(child.bases + [k.value for k in child.keywords]) + else: + stack.extend(child.args.defaults) + stack.extend(d for d in child.args.kw_defaults if d) + if isinstance(child, FUNCTION_NODES): + stack.extend([child.returns] if child.returns else []) + stack.extend( + argument.annotation + for argument in ast.walk(child.args) + if isinstance(argument, ast.arg) and argument.annotation + ) + continue + if isinstance(child, COMPREHENSION_NODES): + # Its loop variables are its own; a walrus binds here. Since 3.12 + # CPython inlines a list, set or dict comprehension (PEP 709; not a + # generator expression), and a nested nonlocal of a loop + # variable's name then names a cell here that only such nonlocals + # set, not the binding further out that 3.11 and a generator + # expression name. Following 3.11 keeps every value either can + # read. + names.update(_comprehension_walrus_targets(child)) + stack.append(child.generators[0].iter) + continue + if isinstance(child, ast.Name) and not isinstance(child.ctx, ast.Load): + names.add(child.id) + elif isinstance(child, (ast.Import, ast.ImportFrom)): + names.update( + (alias.asname or alias.name).split(".")[0] + for alias in child.names + if alias.name != "*" + ) + elif isinstance(child, ast.ExceptHandler) and child.name: + names.add(child.name) + elif isinstance(child, (ast.MatchAs, ast.MatchStar)) and child.name: + names.add(child.name) + elif isinstance(child, ast.MatchMapping) and child.rest: + names.add(child.rest) + stack.extend(ast.iter_child_nodes(child)) + return frozenset(names - declared) + + +def _may_not_run(node, child) -> bool: + """Whether ``child``, a part of ``node``, may not run when ``node`` does: + an operand of ``and`` or ``or`` after the first, a chained comparison's + after the second, a conditional expression's branches, and all of an + ``assert``, which ``python -O`` drops.""" + if isinstance(node, ast.BoolOp): + return child is not node.values[0] + if isinstance(node, ast.Compare): + return any(child is part for part in node.comparators[1:]) + if isinstance(node, ast.IfExp): + return child is not node.test + return isinstance(node, ast.Assert) + + +def _type_parameter_values(node) -> list: + """The bounds and defaults of a definition's type parameters (PEP 695, + Python 3.12+; defaults 3.13+), which Python evaluates when they are + read.""" + values = [] + for parameter in getattr(node, "type_params", []): + for field in ("bound", "default_value"): + value = getattr(parameter, field, None) + if value is not None: + values.append(value) + return values + + +def _alias(scope, value, line): + """What binding ``value`` at ``line`` binds: an _Alias of a name or chain.""" + if value is not None and _dotted(value) is not None: + return _Alias(scope, value, line) + return None + + +def _is_all(node) -> bool: + return isinstance(node, ast.Name) and node.id == "__all__" + + +def _strings(node, current): + """The set of strings ``node`` builds, with ``current`` as __all__.""" + if isinstance(node, (ast.List, ast.Tuple, ast.Set)): + values = [_literal(element) for element in node.elts] + return None if None in values else set(values) + if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Add): + left, right = _strings(node.left, current), _strings(node.right, current) + return None if left is None or right is None else left | right + if _is_all(node): + return set(current) + return None + + +class _Exports(NamedTuple): + # Names __all__ holds once the module has run. + final: frozenset + # Names it held before, or that were removed, which a module importing it + # during an import cycle may see. + earlier: frozenset + + +def _public_names(tree): + """What ``from module import *`` binds, from the module's ``__all__``. + + An _Exports, DEFAULT_PUBLIC without an ``__all__``, or EVERY_NAME for an + ``__all__`` built other than from string literals with ``=``, ``+``, + ``+=``, ``append``, ``extend``, ``insert`` and ``remove`` at the top + level. + """ + uses = {node for node in ast.walk(tree) if _is_all(node)} + if not uses: + return DEFAULT_PUBLIC + names, seen, handled = set(), set(), set() + for statement in tree.body: + value = None + if isinstance(statement, (ast.Assign, ast.AnnAssign)): + targets = ( + statement.targets + if isinstance(statement, ast.Assign) + else [statement.target] + ) + if len(targets) == 1 and _is_all(targets[0]) and statement.value: + value = _strings(statement.value, names) + if value is None: + return EVERY_NAME + names = value + elif ( + isinstance(statement, ast.AugAssign) + and _is_all(statement.target) + and isinstance(statement.op, ast.Add) + ): + value = _strings(statement.value, names) + if value is None: + return EVERY_NAME + names = names | value + elif ( + isinstance(statement, ast.Expr) + and isinstance(statement.value, ast.Call) + and isinstance(statement.value.func, ast.Attribute) + and _is_all(statement.value.func.value) + ): + method, args = statement.value.func.attr, statement.value.args + if method == "append" and len(args) == 1 and _literal(args[0]): + value = {_literal(args[0])} + elif method == "insert" and len(args) == 2 and _literal(args[1]): + value = {_literal(args[1])} + elif method == "extend" and len(args) == 1: + value = _strings(args[0], names) + elif method == "remove": + removed = _literal(args[0]) if len(args) == 1 else None + # Without a literal, any name may be the one removed. + names = names - {removed} if removed else set() + value = set() + if value is None: + return EVERY_NAME + names = names | value + if value is not None: + handled |= {node for node in ast.walk(statement) if _is_all(node)} + seen |= names + if uses - handled: + return EVERY_NAME + return _Exports(frozenset(names), frozenset(seen - names)) + + +class Modules: + """The package's modules, built on demand, for names one imports from another. + + ``sources`` maps module names to (package name, a function returning the + source). A lookup is a pure function of the sources: a cycle is cut + where a lookup meets itself, and a result computed across a cut is not + kept, so no answer depends on the order of lookups. A lookup nested + deeper than ``max_depth`` is cut too, and recorded in ``exhausted``: a + result kept from a shallower lookup could fill that cut in another order. + """ + + def __init__(self, sources: dict, external=None, max_depth=MAX_RESOLUTION_DEPTH): + self.sources = sources + self.external = EXTERNAL_EXPORTS if external is None else external + self.max_depth = max_depth + self.exhausted = set() + self.tops = {} + self.publics = {} + # module -> the package modules it may start, or whose code it may + # reach, when it runs. + self.imported = {} + # module -> the modules it star-imports from. + self.star_bases = {} + # module -> (the package modules it imports or reaches as it runs, + # the packages it star-imports from as it runs). + self.run_imports = {} + # module -> (the package modules code in it that may run later + # imports or reaches, the packages that code star-imports from). + self.later_imports = {} + # package -> the modules below it; every module and package around one. + self._below = None + self._known = None + self._reached = {} + # the packages around a module -> (the package modules that ran + # before it, the package modules their later code may start). + self._ran = {} + self._memo = {} + self._active = set() + self._cuts = 0 + + def top(self, module): + if module not in self.tops: + package, read = self.sources[module] + tree = ast.parse(read(), filename=module) + self.register(module, _build_scopes(tree, module, package, self)[0], tree) + return self.tops[module] + + def register(self, module, top, tree): + if module not in self.tops: + self.tops[module] = top + self.publics[module] = _public_names(tree) + package = self._package(module) + (run, run_stars), (later, later_stars) = _imports_in(tree, package, module) + self._index() + run_reached, later_reached = _references_in( + tree, package, run_stars | later_stars, self._known, module + ) + run, later = run | run_reached, later | later_reached + # Its code may call a function that imports, so every import + # counts. + self.imported[module] = self._imports_of(module, run | later) + self.star_bases[module] = frozenset(run_stars | later_stars) + self.run_imports[module] = ( + self._imports_of(module, run), + frozenset(run_stars), + ) + self.later_imports[module] = ( + self._imports_of(module, later), + frozenset(later_stars), + ) + + def _package(self, module): + return self.sources.get(module, (None,))[0] + + def _imports_of(self, module, names) -> frozenset: + """The package modules ``module`` may start, or reach, through + ``names`` (from ``_imports_in`` and ``_references_in``): each one and + the packages around it. The packages around ``module`` have started + before it. Star imports' submodules are added by ``reached``.""" + started = _started(module) + imported = set() + for name in names: + parts = name.split(".") + for end in range(1, len(parts) + 1): + prefix = ".".join(parts[:end]) + if prefix in self.sources and prefix not in started: + imported.add(prefix) + return frozenset(imported) + + def reached(self, module) -> frozenset: + """The package modules ``module`` may start while it runs, directly + or through those it starts.""" + if module not in self._reached: + seen, stack = set(), [module] + while stack: + current = stack.pop() + self.top(current) + fresh = ( + self.imported[current] + | self.star_loads(current) + | self.called(current) + ) - seen + seen |= fresh + stack.extend(fresh) + self._reached[module] = frozenset(seen) + return self._reached[module] + + def called(self, module) -> frozenset: + """The package modules ``module``'s code may start by calling or + reading code that may run later in the modules that ran before it: + the packages around it, and the package modules their code, as it + ran, imported, directly or through the modules it started.""" + around = _started(module) - {module} + if around not in self._ran: + ran = {package for package in around if package in self.sources} + stack = list(ran) + while stack: + current = stack.pop() + self.top(current) + names, bases = self.run_imports[current] + fresh = (names | self._star_children(bases)) - ran + ran |= fresh + stack.extend(fresh) + starts = set() + for current in ran: + names, bases = self.later_imports[current] + starts |= names | self._star_children(bases) + self._ran[around] = (frozenset(ran), frozenset(starts)) + return self._ran[around][1] - _started(module) + + def star_loads(self, module) -> frozenset: + """The modules ``module``'s star imports of packages may start below + them: what each name a package's ``__all__`` lists, in any version + of it, names below the package, with the packages on the way + (Python imports the package's name, a dot and the listed name, + dotted or not), or every module below it where ``__all__`` is built + otherwise. A star import of a package without ``__all__`` starts + none. The packages around ``module`` have started before it.""" + return self._star_children(self.star_bases[module]) - _started(module) + + def _index(self): + if self._below is None: + self._below, known = {}, set() + for name in self.sources: + known |= _started(name) + parts = name.split(".") + for end in range(1, len(parts)): + self._below.setdefault(".".join(parts[:end]), set()).add(name) + self._known = frozenset(known) + + def _star_children(self, bases) -> frozenset: + """The modules below them star imports from ``bases`` may start.""" + self._index() + loads = set() + for base in bases: + if base not in self.sources: + continue + self.top(base) + public = self.publics[base] + if public == EVERY_NAME: + loads |= self._below.get(base, set()) + elif isinstance(public, _Exports): + for listed in public.final | public.earlier: + parts = f"{base}.{listed}".split(".") + for end in range(base.count(".") + 2, len(parts) + 1): + path = ".".join(parts[:end]) + if path in self.sources: + loads.add(path) + return frozenset(loads) + + def cyclic(self, importer, base) -> bool: + """Whether ``importer``'s star import of ``base`` may run while + ``base`` is still running, and see only what it has bound so far: + when ``base`` may start ``importer``, or is it.""" + return importer == base or importer in self.reached(base) + + def guarded(self, key, cut, compute): + """``compute()``, or ``cut`` where the lookup meets itself.""" + if key in self._memo: + return self._memo[key] + if key in self._active or len(self._active) > self.max_depth: + if key not in self._active: + self.exhausted.add(key) + self._cuts += 1 + return cut + self._active.add(key) + cuts = self._cuts + try: + result = compute() + finally: + self._active.discard(key) + if self._cuts == cuts and not isinstance(key, _Alias): + self._memo[key] = result + return result + + def binding(self, module, name): + """(bound, targets, falls through) for ``module``'s top-level + ``name`` once it ran: whether any statement binds it, what to, and + whether it may be unbound after all.""" + + def compute(): + targets, bound, falls = self.top(module).own(name, None, True) + return bound, frozenset(targets), falls + + cut = (False, frozenset(), True) + return self.guarded(("binding", module, name), cut, compute) + + def expand(self, target, start=1) -> frozenset: + """``target`` and what else it names, where a prefix of it at least + ``start`` parts long is a package module or one ``external`` maps and + the next part is a name that module binds. + + The binding's targets, expanded, replace that prefix and name, and + only the rest is expanded further: a dotted name is both a module + path and a chain of attributes, and ``a.b`` may bind ``b`` to a path + through the submodule ``a.b``. + """ + + def compute(): + found = {target} + parts = target.split(".") + for index in range(start, len(parts)): + base, name = ".".join(parts[:index]), parts[index] + if base in self.sources: + bound = self.binding(base, name)[1] + else: + mapped = self.external.get(base, {}).get(name) + bound = {mapped} if mapped else set() + for head in self.expand_all(bound): + named = ".".join([head] + parts[index + 1 :]) + if named != target: + found |= self.expand(named, len(head.split("."))) + return frozenset(found) + + key = ("expand", target, start) + return self.guarded(key, frozenset({target}), compute) + + def expand_all(self, targets) -> set: + return set().union(*map(self.expand, targets)) + + def star_binds(self, base, name, importer): + """What ``from base import *`` in ``importer`` binds ``name`` to, or + None.""" + if base not in self.sources: + mapped = self.external.get(base, {}).get(name) + return {mapped} if mapped else None + self.top(base) + public = self.publics[base] + if public == DEFAULT_PUBLIC: + exported = not name.startswith("_") + elif public == EVERY_NAME: + exported = True + else: + exported = name in public.final or name in public.earlier + if not exported and not name.startswith("_"): + # In an import cycle it may run before base sets __all__, + # and then binds every public name base has bound. + exported = self.cyclic(importer, base) and self.binding(base, name)[0] + if not exported: + return None + bound, targets, _ = self.binding(base, name) + if bound: + return set(targets) + if isinstance(public, _Exports): + return {f"{base}.{name}"} # A submodule __all__ lists. + return None + + def star_binds_surely(self, base, name, importer) -> bool: + """Whether ``from base import *`` in ``importer``, which star_binds + says may bind ``name``, is sure to.""" + if base not in self.sources: + return True + if self.cyclic(importer, base): + return False # It may run before base binds anything. + public = self.publics[base] + if public == EVERY_NAME: + return False + if isinstance(public, _Exports) and name not in public.final: + return False # Only an earlier or removed __all__ held it. + bound, _, falls = self.binding(base, name) + return not (bound and falls) + + +def _build_scopes(tree, module, package, modules): + """(module scope, the innermost scope of every node, the (first line, + last line) of the innermost statement around every node) of ``tree``.""" + top = _Scope("module", "", None, modules) + scope_of, statement_of = {}, {} + + def assign(target, value, scope, conditional, evaluated_in=None): + """Bind ``target`` in ``scope`` to ``value``, evaluated in + ``evaluated_in`` (by default ``scope``).""" + scope_of[target] = scope + if isinstance(target, ast.Name): + line = target.lineno + alias = _alias(evaluated_in or scope, value, line) + scope.bind(target.id, line, alias, conditional) + elif ( + isinstance(target, (ast.Tuple, ast.List)) + and isinstance(value, (ast.Tuple, ast.List)) + and len(target.elts) == len(value.elts) + ): + for element, element_value in zip(target.elts, value.elts): + assign(element, element_value, scope, conditional, evaluated_in) + else: + visit(target, scope, conditional) + + def bind_arguments(inner, outer, arguments, line): + positional = arguments.posonlyargs + arguments.args + defaults = [None] * (len(positional) - len(arguments.defaults)) + defaults += arguments.defaults + pairs = list(zip(positional, defaults)) + pairs += list(zip(arguments.kwonlyargs, arguments.kw_defaults)) + pairs += [(arguments.vararg, None), (arguments.kwarg, None)] + for argument, default in pairs: + if argument is not None: + inner.bind(argument.arg, line, _alias(outer, default, line)) + + def declare(inner, body): + """Send names a global or nonlocal statement in ``body`` names to + the scope that binds them.""" + stack = list(body) + while stack: + node = stack.pop() + if isinstance(node, ast.Global): + for name in node.names: + inner.declared[name] = top + elif isinstance(node, ast.Nonlocal): + # The nearest enclosing function that makes the name local, + # wherever in its body it binds it; class bodies are skipped. + for name in node.names: + owner = inner.parent + while owner is not None and not ( + owner.kind in ("function", "lambda") + and name in owner.local_names + ): + owner = owner.parent + if owner is not None: + inner.declared[name] = owner + elif not isinstance(node, FUNCTION_NODES + (ast.ClassDef, ast.Lambda)): + stack.extend(ast.iter_child_nodes(node)) + + def visit_annotations(annotations, scope, line, conditional): + # Evaluated in the outer scope, with the statement or later (see + # _Scope.resolve); the statement's own bindings are not sure for + # them, as for any use in it. + if not annotations: + return + inner = _Scope("annotation", scope.qualname, scope, modules, line) + for annotation in annotations: + visit(annotation, inner, conditional) + + def visit_function(node, scope, conditional): + # Decorators and defaults are evaluated in the outer scope. + outer = list(getattr(node, "decorator_list", [])) + outer += node.args.defaults + [d for d in node.args.kw_defaults if d] + for child in outer: + visit(child, scope, conditional) + if not isinstance(node, ast.Lambda): + annotations = [node.returns] if node.returns else [] + annotations += [ + argument.annotation + for argument in ast.walk(node.args) + if isinstance(argument, ast.arg) and argument.annotation + ] + annotations += _type_parameter_values(node) + visit_annotations(annotations, scope, node.lineno, conditional) + name = getattr(node, "name", "") + if not isinstance(node, ast.Lambda): + top_level = scope.parent is None + target = f"{module}.{name}" if top_level else None + scope.bind(name, node.lineno, target, conditional) + qualname = name if scope.parent is None else f"{scope.qualname}.{name}" + kind = "lambda" if isinstance(node, ast.Lambda) else "function" + inner = _Scope(kind, qualname, scope, modules, node.lineno) + if getattr(node, "decorator_list", None): + # A decorator may call the function before its name is bound. A + # lambda's body is part of its own statement, which every use + # in it already treats as unfinished (resolve). + inner.unfinished = statement_span[-1] + inner.local_names = _function_locals(node) + body = node.body if isinstance(node.body, list) else [node.body] + declare(inner, body) + bind_arguments(inner, scope, node.args, node.lineno) + for child in body: + visit(child, inner) + + def visit_class(node, scope, conditional): + for child in node.decorator_list + node.bases: + visit(child, scope, conditional) + for keyword in node.keywords: + visit(keyword.value, scope, conditional) + visit_annotations(_type_parameter_values(node), scope, node.lineno, conditional) + top_level = scope.parent is None + target = f"{module}.{node.name}" if top_level else None + scope.bind(node.name, node.lineno, target, conditional) + qualname = node.name if top_level else f"{scope.qualname}.{node.name}" + inner = _Scope("class", qualname, scope, modules, node.lineno) + inner.unfinished = statement_span[-1] + inner.local_names = _block_locals(node.body) + declare(inner, node.body) + for child in node.body: + visit(child, inner) + + def visit_comprehension(node, scope, conditional): + # The first iterable is evaluated in the enclosing scope; the rest + # in the comprehension's own, reported under the enclosing name. + visit(node.generators[0].iter, scope, conditional) + # A generator expression runs the rest when it is iterated. + kind = "generator" if isinstance(node, ast.GeneratorExp) else "comprehension" + inner = _Scope(kind, scope.qualname, scope, modules, node.lineno) + # The rest may run any number of times, a generator expression's + # only as it is iterated, so a walrus in it may not bind. + for index, generator in enumerate(node.generators): + scope_of[generator] = inner + if index: + visit(generator.iter, inner, True) + assign(generator.target, None, inner, False) + for condition in generator.ifs: + visit(condition, inner, True) + if isinstance(node, ast.DictComp): + elements = [node.key, node.value] + else: + elements = [node.elt] + for element in elements: + visit(element, inner, True) + + # (first line, last line) of each statement being visited, innermost + # last; a definition's decorators are part of its statement. + statement_span = [(1, 1)] + + def visit(node, scope, conditional=False): + if not isinstance(node, ast.stmt): + return visit_node(node, scope, conditional) + lines = [node.lineno] + lines += [d.lineno for d in getattr(node, "decorator_list", [])] + statement_span.append((min(lines), node.end_lineno)) + try: + return visit_node(node, scope, conditional) + finally: + statement_span.pop() + + def visit_node(node, scope, conditional): + scope_of[node] = scope + statement_of[node] = statement_span[-1] + if isinstance(node, FUNCTION_NODES + (ast.Lambda,)): + return visit_function(node, scope, conditional) + if isinstance(node, ast.ClassDef): + return visit_class(node, scope, conditional) + if isinstance(node, COMPREHENSION_NODES): + return visit_comprehension(node, scope, conditional) + if isinstance(node, TYPE_ALIAS): + # Its value is evaluated when it is read, like an annotation. + scope_of[node.name] = scope + scope.bind(node.name.id, node.lineno, None, conditional) + values = [node.value] + _type_parameter_values(node) + return visit_annotations(values, scope, node.lineno, conditional) + if isinstance(node, ast.Assign): + visit(node.value, scope, conditional) + for target in node.targets: + assign(target, node.value, scope, conditional) + return + if isinstance(node, ast.AnnAssign): + body = scope.kind in ("module", "class") + if body and node.simple: + visit_annotations([node.annotation], scope, node.lineno, conditional) + else: + # 3.14 never evaluates this one; earlier versions do, with + # the statement, outside a function. + visit(node.annotation, scope, conditional) + if node.value is not None: + visit(node.value, scope, conditional) + elif body and isinstance(node.target, ast.Name): + # An annotation alone binds no value here; in a function it + # still makes the name local. + scope_of[node.target] = scope + return + return assign(node.target, node.value, scope, conditional) + if isinstance(node, ast.NamedExpr): + visit(node.value, scope, conditional) + # The target belongs to the scope around any comprehensions or + # annotations; the value is evaluated where the expression is. + owner = scope + while owner.kind in ("comprehension", "generator", "annotation"): + owner = owner.parent + return assign(node.target, node.value, owner, conditional, scope) + if isinstance(node, ast.Import): + for alias in node.names: + if alias.asname: + scope.bind(alias.asname, node.lineno, alias.name, conditional) + else: + root = alias.name.split(".")[0] + scope.bind(root, node.lineno, root, conditional) + elif isinstance(node, ast.ImportFrom): + base = _absolute(node.module or "", node.level, package) + for alias in node.names: + if alias.name == "*": + scope.stars.append((node.lineno, base, conditional, module)) + else: + # A submodule or a name base binds: expand decides. + name = alias.asname or alias.name + scope.bind(name, node.lineno, f"{base}.{alias.name}", conditional) + elif isinstance(node, ast.Name) and isinstance(node.ctx, ast.Store): + scope.bind(node.id, node.lineno, None, conditional) + elif isinstance(node, ast.Name) and isinstance(node.ctx, ast.Del): + scope.delete(node.id, node.lineno) + elif isinstance(node, ast.ExceptHandler) and node.name: + # Python deletes the name when the handler ends. + scope.bind(node.name, node.lineno, None, True) + scope.delete(node.name, node.end_lineno) + inside = conditional or isinstance(node, COMPOUND_NODES) + for child in ast.iter_child_nodes(node): + visit(child, scope, inside or _may_not_run(node, child)) + + visit(tree, top) + return top, scope_of, statement_of + + +def _reads_an_id(node) -> bool: + for child in ast.walk(node): + if isinstance(child, ast.Name) and child.id.endswith("_id"): + return True + if ( + isinstance(child, ast.Constant) + and isinstance(child.value, str) + and child.value.endswith("_id") + ): + return True + return False + + +def _is_xorshift(left, right, shift_of) -> bool: + """``x ^ (x >> k)`` or ``x ^ (x << k)``, in either order. + + ``shift_of`` gives the shifted operand of a shift, or None. + """ + for shifted, other in ((left, right), (right, left)): + operand = shift_of(shifted) + if operand is not None and ast.unparse(operand) == ast.unparse(other): + return True + return False + + +def _argument(call, position, keyword): + """The argument ``call`` passes at ``position`` or as ``keyword``.""" + if len(call.args) > position: + return call.args[position] + for passed in call.keywords: + if passed.arg == keyword: + return passed.value + return None + + +def _anchor(node, module): + """The package a literal, or the module's ``__name__``, names, else None + (callers read anything computed, ``__package__`` or ``__spec__.parent`` + too, as the module's own package).""" + if _literal(node) is not None: + return _literal(node) + if isinstance(node, ast.Name) and node.id == "__name__": + return module + return None + + +def _is_none(node) -> bool: + return isinstance(node, ast.Constant) and node.value is None + + +def _dict_entries(node) -> list: + """(key, value) for each entry of dict literal ``node``, in order, a + literal ``**`` dict's own entries in its place; a computed ``**`` entry + gives (None, the mapping).""" + entries = [] + for key, value in zip(node.keys, node.values): + if key is None and isinstance(value, ast.Dict): + entries += _dict_entries(value) + else: + entries.append((key, value)) + return entries + + +def _iterated(node) -> list: + """The items Python gets iterating ``node``, as far as it is literal: a + string's characters, a list's, tuple's or set's elements and a dict's + keys, a ``*`` or ``**`` entry's own in its place. An item that is not a + literal string is None; a computed iterable gives none.""" + if isinstance(node, (ast.List, ast.Tuple, ast.Set)): + items = [] + for element in node.elts: + if isinstance(element, ast.Starred): + items += _iterated(element.value) + else: + items.append(_literal(element)) + return items + if isinstance(node, ast.Dict): + return [_literal(key) for key, _ in _dict_entries(node)] + return list(_literal(node) or "") + + +def _globals_package(node, package, module): + """The package a relative ``__import__`` with ``globals`` ``node`` + imports from, as Python reads it: ``__package__``, else + ``__spec__.parent``, else ``__name__``, the package around it unless + ``__path__`` is set; None where it has neither, and Python raises. + A dict literal's literal ``**`` entries are read in place; anything + else (``globals()``, a computed mapping, or a dict literal with a + computed key or ``**`` entry) is read as the module's own globals.""" + if not isinstance(node, ast.Dict): + return package + pairs = _dict_entries(node) + if not all(isinstance(key, ast.Constant) for key, _ in pairs): + return package + # The last of a repeated key wins, as in the dict Python builds; a + # None value then reads as absent. Keys that aren't strings are + # never ones Python reads. + entries = {key.value: value for key, value in pairs} + given = {key: value for key, value in entries.items() if not _is_none(value)} + if "__package__" in given: + return _anchor(given["__package__"], module) or package + if "__spec__" in given: + return package + if "__name__" not in given: + return None + name = _anchor(given["__name__"], module) + if name is None: + return package + return name if "__path__" in entries else name.rpartition(".")[0] + + +def _dynamic_imports(call, called, package, module=None) -> list: + """Absolute names a call of a dynamic import function imports, where its + arguments are literals: ``base.*`` where an ``__import__`` ``fromlist`` + holds ``*``.""" + names = [] + if "importlib.import_module" in called: + name = _literal(_argument(call, 0, "name")) + given = _argument(call, 1, "package") + relative_to = _anchor(given, module) + if relative_to is None and given is not None and _literal(given) is None: + # A computed package is read as the module's own; without one, + # Python raises. + relative_to = None if _is_none(given) else package + if name and name.startswith("."): + level = len(name) - len(name.lstrip(".")) + names.append(_absolute(name.lstrip("."), level, relative_to)) + elif name: + names.append(name) + if "builtins.__import__" in called: + name = _literal(_argument(call, 0, "name")) + level = _argument(call, 4, "level") + level = level.value if isinstance(level, ast.Constant) else 0 + if name is not None and isinstance(level, int): + relative_to = _globals_package( + _argument(call, 1, "globals"), package, module + ) + base = _absolute(name, level, relative_to) if level else name + names.append(base) + items = _iterated(_argument(call, 3, "fromlist")) + names += [f"{base}.{item}" for item in items if item] + return names + + +# Code that may run after the statement holding it, when it is called or +# read, so after the module has run: as may annotations (deferred from 3.14), +# type parameters and ``type`` values. +LATER_NODES = (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda, ast.GeneratorExp) +LATER_FIELDS = ("annotation", "type_params") + + +def _later_walk(tree): + """(node, whether it may run later) for every node of ``tree``: in a + function, lambda or generator expression (with its decorators, defaults + and first iterable, which run at once, to be safe), an annotation, type + parameters or a ``type`` alias.""" + stack = [(tree, False)] + while stack: + node, later = stack.pop() + yield node, later + later = later or isinstance(node, LATER_NODES) or isinstance(node, TYPE_ALIAS) + for field, value in ast.iter_fields(node): + for child in value if isinstance(value, list) else [value]: + if isinstance(child, ast.AST): + stack.append((child, later or field in LATER_FIELDS)) + + +def _imports_in(tree, package, module=None): + """((names, star bases) of the imports that run with ``tree``'s module, + (names, star bases) of those in code that may run later (``_later_walk``). + + Names are absolute, as an import, or a call whose literal arguments name + a module as ``importlib.import_module`` or ``__import__`` would read + them, imports them, whatever name the call is made under: an alias, a + re-export or ``getattr`` may make it one. Star bases are the packages a + star import, or ``__import__`` with ``*`` in its ``fromlist``, loads + from.""" + found = {False: (set(), set()), True: (set(), set())} + for node, later in _later_walk(tree): + names, stars = found[later] + if isinstance(node, ast.Import): + names.update(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom): + base = _absolute(node.module or "", node.level, package) + names.add(base) + names.update(f"{base}.{alias.name}" for alias in node.names) + if any(alias.name == "*" for alias in node.names): + stars.add(base) + elif isinstance(node, ast.Call): + for name in _dynamic_imports(node, DYNAMIC_IMPORTS, package, module): + names.add(name) + if name.endswith(".*"): + stars.add(name[: -len(".*")]) + return found[False], found[True] + + +# What _references_in's rules read, with the paths on the way. +REFERENCE_NAMES = frozenset( + {"getattr", "builtins", "builtins.getattr", "sys", "sys.modules", "sys.modules.get"} +) + + +def _references_in(tree, package, stars, known, module=None): + """(paths code that runs with ``tree``'s module may name, paths code that + may run later may name; ``_later_walk``). + + Another module may have loaded a module, which is then an attribute of + its package, so code can reach it, and call its functions, without + importing it. Each name, attribute chain, ``getattr`` with a literal + name, ``sys.modules`` read with a literal key (``[]`` or ``get``) and + literal dynamic import (the module or its top package, which it + returns), and each branch of a conditional expression, ``and`` or + ``or``, names every path its first name may: what any import in the + module, or assignment of such an expression (or walrus), binds the name + to (an assignment may name the name itself, as in ``x = x.sub``), the name + itself, the name below each package in ``stars`` (a star import of a + package binds its loaded submodules), and in a package's ``__init__`` + the name below the package. Only paths in ``known`` (each package module + and the packages around it), from which a module may yet be named, and + those the rules read are kept. Which paths are package modules is the + caller's to say.""" + bound = {} + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + name = alias.asname or alias.name.split(".")[0] + bound.setdefault(name, []).append(alias.name if alias.asname else name) + elif isinstance(node, ast.ImportFrom): + base = _absolute(node.module or "", node.level, package) + for alias in node.names: + if alias.name != "*": + name = alias.asname or alias.name + bound.setdefault(name, []).append(f"{base}.{alias.name}") + elif isinstance(node, (ast.Assign, ast.AnnAssign, ast.NamedExpr)): + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + for target in targets: + if isinstance(target, ast.Name) and node.value is not None: + bound.setdefault(target.id, []).append(node.value) + around = list(stars) + ([package] if package and module in (None, package) else []) + keep = known | REFERENCE_NAMES + names = {} + + def name_paths(name): + if name not in names: + names[name] = ({name} | {f"{base}.{name}" for base in around}) & keep + return names[name] + + def paths(node): + return _paths(node) & keep + + def _paths(node): + if isinstance(node, ast.Name): + return name_paths(node.id) + if isinstance(node, ast.Attribute): + return {f"{path}.{node.attr}" for path in paths(node.value)} + if isinstance(node, ast.Subscript): + key = _literal(node.slice) + if key and "sys.modules" in paths(node.value): + return {key} + return set() + if isinstance(node, ast.IfExp): + return paths(node.body) | paths(node.orelse) + if isinstance(node, ast.BoolOp): + return set().union(*map(paths, node.values)) + if isinstance(node, ast.NamedExpr): + return paths(node.value) + if not isinstance(node, ast.Call): + return set() + called = paths(node.func) + if called & {"getattr", "builtins.getattr"} and len(node.args) >= 2: + attribute = _literal(node.args[1]) + if attribute: + return {f"{path}.{attribute}" for path in paths(node.args[0])} + if "sys.modules.get" in called and node.args and _literal(node.args[0]): + return {_literal(node.args[0])} + returned = set() + for name in _dynamic_imports(node, DYNAMIC_IMPORTS, package, module): + if not name.endswith(".*"): + returned |= {name, name.split(".")[0]} + return returned + + # Bind each name to what its imports and assignments name, until no + # path is new: kept paths are finite, so this ends. + changed = True + while changed: + changed = False + for name, values in bound.items(): + for value in values: + fresh = ({value} & keep if isinstance(value, str) else paths(value)) - ( + name_paths(name) + ) + if fresh: + names[name] |= fresh + changed = True + + found = {False: set(), True: set()} + for node, later in _later_walk(tree): + if isinstance(node, (ast.Name, ast.Attribute, ast.Subscript, ast.Call)): + found[later] |= paths(node) + return found[False], found[True] + + +def _started(module) -> frozenset: + """``module`` and the packages around it, which start before it runs.""" + parts = module.split(".") + return frozenset(".".join(parts[:end]) for end in range(1, len(parts) + 1)) + + +class Analysis(NamedTuple): + module: str + # (line, enclosing scope, kind), without draw helper calls. + findings: list + # (line, enclosing scope, targets) for each use that may name something + # in the package. + references: list + # Names of the module's top-level functions, also those defined in an + # if or try. + functions: list + # (line, absolute module name) for each import, static or literal dynamic. + imports: list + + +def analyse( + source: str, + filename: str = "", + module: str = "m", + package: str | None = None, + modules: Modules | None = None, +) -> Analysis: + """Draw patterns, package references and imports of one module. + + ``modules`` resolves names imported from other package modules; without + it the module stands alone. + """ + tree = ast.parse(source, filename=filename) + if modules is None: + modules = Modules({module: (package, lambda: source)}) + top, scope_of, statement_of = _build_scopes(tree, module, package, modules) + modules.register(module, top, tree) + parents = { + child: node for node in ast.walk(tree) for child in ast.iter_child_nodes(node) + } + root_package = module.split(".")[0] + "." + findings, references, imports = [], [], [] + + def scope_for(node): + while node not in scope_of: + node = parents[node] + return scope_of[node] + + def statement_for(node): + while node not in statement_of: + node = parents[node] + return statement_of[node] + + def report(node, kind): + findings.append((node.lineno, scope_for(node).qualname, kind)) + + def report_targets(node, targets): + for kind in sorted({_classify(target) for target in targets} - {None}): + report(node, kind) + + def targets_of(node) -> set: + dotted = _dotted(node) + if dotted is None: + return set() + root, attributes = dotted + return modules.expand_all( + ".".join([target] + attributes) + for target in scope_for(node).resolve( + root, node.lineno, statement=statement_for(node) + ) + ) + + def shift_of(node): + if isinstance(node, ast.BinOp) and isinstance( + node.op, (ast.RShift, ast.LShift) + ): + return node.left + if isinstance(node, ast.Call) and node.args: + if targets_of(node.func) & SHIFT_FUNCTIONS: + return node.args[0] + return None + + def check_import(node, name): + imports.append((node.lineno, name)) + if _is_stochastic(name): + report(node, "stochastic import") + elif name.split(".")[0] in HASH_MODULES: + report(node, "hash library") + + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + check_import(node, alias.name) + elif isinstance(node, ast.ImportFrom): + base = _absolute(node.module or "", node.level, package) + names = [base] + [f"{base}.{alias.name}" for alias in node.names] + imports.extend((node.lineno, name) for name in names) + if any(map(_is_stochastic, names)): + report(node, "stochastic import") + elif base.split(".")[0] in HASH_MODULES: + report(node, "hash library") + elif isinstance(node, (ast.Name, ast.Attribute)) and isinstance( + node.ctx, ast.Load + ): + parent = parents.get(node) + if isinstance(parent, ast.Attribute) and parent.value is node: + continue # Only the whole chain counts. + targets = targets_of(node) + report_targets(node, targets) + named = frozenset(t for t in targets if t.startswith(root_package)) + if named: + references.append((node.lineno, scope_for(node).qualname, named)) + if isinstance(node, ast.Call): + func = node.func + called = targets_of(func) + if called & DYNAMIC_IMPORTS: + for name in _dynamic_imports(node, called, package, module): + check_import(node, name) + if "builtins.getattr" in called and len(node.args) >= 2: + attribute = _literal(node.args[1]) + if attribute: + owners = targets_of(node.args[0]) + report_targets( + node, + modules.expand_all(f"{owner}.{attribute}" for owner in owners), + ) + if called & XOR_FUNCTIONS and len(node.args) == 2: + if _is_xorshift(*node.args, shift_of): + report(node, "xorshift") + if called & MOD_FUNCTIONS and node.args and _reads_an_id(node.args[0]): + report(node, "id modulo") + if isinstance(func, ast.Attribute): + keywords = {keyword.arg for keyword in node.keywords} + if func.attr in SAMPLING_METHODS or ( + func.attr == "sample" + and (node.args or keywords & SAMPLING_KEYWORDS) + ): + report(node, "sampling call") + elif isinstance(node, ast.Constant): + value = node.value + if ( + isinstance(value, int) + and not isinstance(value, bool) + and value in HASH_CONSTANTS + ): + report(node, "hash constant") + elif isinstance(node, (ast.BinOp, ast.AugAssign)): + left = node.left if isinstance(node, ast.BinOp) else node.target + right = node.right if isinstance(node, ast.BinOp) else node.value + if isinstance(node.op, ast.BitXor) and _is_xorshift(left, right, shift_of): + report(node, "xorshift") + if isinstance(node.op, ast.Mod) and _reads_an_id(left): + report(node, "id modulo") + functions = [ + node.name + for node in ast.walk(tree) + if isinstance(node, FUNCTION_NODES) and scope_of.get(node) is top + ] + return Analysis(module, findings, references, functions, imports) + + +def helper_calls(analysis: Analysis, helpers) -> list: + """(line, scope, "draw helper call") for each use of a draw helper.""" + return [ + (line, scope, "draw helper call") + for line, scope, targets in analysis.references + if any(_matches(target, helpers) for target in targets) + ] + + +def draw_helpers(analyses: dict, exempt: dict) -> set: + """Top-level functions with a finding ``exempt`` does not cover. + + A function that uses one of them is one too, so this is a fixed point + over the package. + """ + helpers = set() + while True: + added = set() + for relative, analysis in analyses.items(): + allowed = exempt.get(relative, {}) + counted = Counter( + (scope, kind) + for _, scope, kind in analysis.findings + + helper_calls(analysis, helpers) + ) + for name in analysis.functions: + target = f"{analysis.module}.{name}" + if target not in helpers and any( + count > allowed.get(key, 0) + for key, count in counted.items() + if key[0] == name or key[0].startswith(name + ".") + ): + added.add(target) + if not added: + return helpers + helpers |= added + + +def draw_findings(source: str, filename: str = "", exempt=None) -> list: + """(line, enclosing scope, kind) for every draw pattern in ``source``. + + ``exempt`` is the file's EXEMPT entries, which decide which of its + functions are draw helpers. + """ + analysis = analyse(source, filename) + helpers = draw_helpers({filename: analysis}, {filename: exempt or {}}) + return sorted(analysis.findings + helper_calls(analysis, helpers)) + + +def count_findings(findings_by_file: dict) -> dict: + """{file: {(scope, kind): [lines]}} from {file: draw_findings(...)}.""" + counted = {} + for relative, findings in findings_by_file.items(): + for line, scope, kind in findings: + counted.setdefault(relative, {}).setdefault((scope, kind), []).append(line) + return counted + + +def check_allowlists(counted: dict, ratchet: dict, exempt: dict): + """(violations, stale entries) of counted findings against the lists.""" + violations, stale = [], [] + allowed = {} + for allowlist in (ratchet, exempt): + for relative, entries in allowlist.items(): + for key, count in entries.items(): + allowed[(relative, key)] = allowed.get((relative, key), 0) + count + for relative, entries in sorted(counted.items()): + for (scope, kind), lines in sorted(entries.items()): + limit = allowed.get((relative, (scope, kind)), 0) + if len(lines) > limit: + violations.append( + f"{relative}:{','.join(map(str, lines))} in {scope}: {kind} " + f"({len(lines)} found, {limit} allowed)" + ) + for (relative, (scope, kind)), limit in sorted(allowed.items()): + found = len(counted.get(relative, {}).get((scope, kind), [])) + if found < limit: + stale.append( + f"{relative} in {scope}: {kind} ({limit} allowed, {found} found)" + ) + return violations, stale + + +def package_modules(): + """(relative path, module name, package name, path) of every module.""" + for path in sorted(PACKAGE.rglob("*.py")): + if TESTS in path.parents: + continue + relative = path.relative_to(PACKAGE).as_posix() + dotted = f"{PACKAGE.name}.{relative.removesuffix('.py').replace('/', '.')}" + if dotted.endswith(".__init__"): + module = dotted.removesuffix(".__init__") + yield relative, module, module, path + else: + yield relative, dotted, dotted.rsplit(".", 1)[0], path + + +def _reader(path): + return lambda: path.read_text(encoding="utf-8") + + +class PackageScan(NamedTuple): + counts: dict + helpers: set + imports: dict + modules: Modules + + +@cache +def package_scan() -> PackageScan: + """Every module analysed once per test session.""" + listed = list(package_modules()) + modules = Modules( + {module: (package, _reader(path)) for _, module, package, path in listed} + ) + analyses = { + relative: analyse( + path.read_text(encoding="utf-8"), relative, module, package, modules + ) + for relative, module, package, path in listed + } + helpers = draw_helpers(analyses, EXEMPT) + counts = count_findings( + { + relative: analysis.findings + helper_calls(analysis, helpers) + for relative, analysis in analyses.items() + } + ) + imports = {relative: analysis.imports for relative, analysis in analyses.items()} + return PackageScan(counts, helpers, imports, modules) + + +@pytest.fixture(scope="module", autouse=True) +def _release_package_scan(): + """Free the package scan when this module's tests finish.""" + yield + package_scan.cache_clear() + + +def _kinds(source: str) -> set: + return {kind for _, _, kind in draw_findings(source)} + + +def _on_this_python(cases, minor=-1) -> list: + """The cases whose syntax this Python parses. + + Each case gives the Python 3 minor version it needs at index ``minor``. + Filtering them as the tests are collected, rather than skipping them, + leaves no skipped outcome on an older Python. + """ + return [case for case in cases if sys.version_info >= (3, case[minor])] + + +@pytest.mark.parametrize( + "source, kind", + [ + ("import random\n", "stochastic import"), + ("from random import uniform\nuniform()\n", "random source"), + ("import random as rnd\nrnd.random()\n", "random source"), + ("import secrets\n", "stochastic import"), + ("import numpy.random as npr\nnpr.uniform()\n", "random source"), + ("from numpy import random as r\nr.random(3)\n", "random source"), + ("from numpy import random as r\n", "stochastic import"), + ("from numpy.random import default_rng as g\ng(43)\n", "random source"), + ("np.random.seed(1)\n", "random source"), + ("import numpy as n\nn.random.default_rng(0)\n", "random source"), + ("def f():\n return np.random.random(3)\n", "random source"), + ("import uuid\nuuid.uuid4()\n", "random source"), + ("import os\nos.urandom(8)\n", "random source"), + # Aliases, star imports and scopes. + ("import numpy as np\nxp = np\nxp.random.uniform()\n", "random source"), + ("r = np.random\ndef f():\n return r.uniform()\n", "random source"), + ("xp, k = np, 1\nxp.random.uniform()\n", "random source"), + ("def f(xp=np):\n return xp.random.uniform()\n", "random source"), + ("from numpy import *\nrandom.uniform()\n", "random source"), + ("getattr(np, 'random').uniform()\n", "random source"), + ( + "import numpy as np\n" + "def f():\n" + " return [np for np in np.random.uniform(size=2)]\n", + "random source", + ), + ("if (r := np.random):\n r.uniform()\n", "random source"), + # A walrus in a comprehension binds outside it, to a value read inside. + ( + "import numpy as np\n[(lib := np) for _ in [0]]\nlib.random.uniform()\n", + "random source", + ), + # A name may refer to any binding that may be the one read. + ( + "from policyengine_core.model_api import np as lib\nlib.random.uniform()\n", + "random source", + ), + # An attribute of a module is what the module binds under the name. + ( + "import policyengine_core.model_api as api\napi.np.random.uniform()\n", + "random source", + ), + ( + "from policyengine_core import model_api\nmodel_api.np.random.uniform()\n", + "random source", + ), + ( + "import policyengine_core.model_api as api\nlib = api.np\nlib.random.uniform()\n", + "random source", + ), + ( + "import numpy as lib\n" + "values = [lib.random.uniform() for _ in range(1)]\n" + "lib = object()\n", + "random source", + ), + ( + "import numpy as np\n" + "class Outer:\n" + " np = object()\n" + " class Inner:\n" + " value = np.random.uniform()\n", + "random source", + ), + ( + "class Outer:\n hash = 1\n class Inner:\n key = hash('x')\n", + "builtin hash", + ), + ("import numpy as lib; alias = lib; alias.random.uniform()\n", "random source"), + ( + "try:\n" + " import numpy as lib\n" + "except ImportError:\n" + " lib = None\n" + "lib.random.uniform()\n", + "random source", + ), + ( + "def setup():\n" + " global lib\n" + " import numpy as lib\n" + "def f():\n" + " return lib.random.uniform()\n", + "random source", + ), + ( + "def outer():\n" + " lib = None\n" + " def setup():\n" + " nonlocal lib\n" + " import numpy as lib\n" + " return lib.random.uniform()\n", + "random source", + ), + # A global name skips the functions between it and the module. + ( + "import numpy as lib\n" + "def outer():\n" + " lib = None\n" + " def inner():\n" + " global lib\n" + " return lib.random.random()\n" + " return inner\n", + "random source", + ), + # nonlocal names the nearest function that binds the name, not the + # nearest function, and wherever in its body the binding is. + ( + "def outer():\n" + " import numpy as lib\n" + " def middle():\n" + " def inner():\n" + " nonlocal lib\n" + " lib.random.random()\n" + " lib = None\n" + " inner()\n" + " middle()\n", + "random source", + ), + ( + "def outer():\n" + " def middle():\n" + " def inner():\n" + " nonlocal lib\n" + " lib = None\n" + " return lib.random.random()\n" + " return inner\n" + " import numpy as lib\n" + " return middle\n", + "random source", + ), + # A walrus in a nested function's annotation binds in the function + # around it, so middle binds lib and inner reads numpy. + ( + "import numpy as np\n" + "def outer():\n" + " lib = None\n" + " def middle():\n" + " def annotated(x: (lib := np)):\n" + " pass\n" + " def inner():\n" + " nonlocal lib\n" + " return lib.random.random()\n" + " return inner\n" + " return middle\n", + "random source", + ), + ( + "import numpy as np\n" + "def outer():\n" + " lib = None\n" + " def middle():\n" + " def annotated() -> (lib := np):\n" + " pass\n" + " def inner():\n" + " nonlocal lib\n" + " return lib.random.random()\n" + " return inner\n" + " return middle\n", + "random source", + ), + # A walrus in a lambda binds in the lambda, even inside a + # comprehension, so middle doesn't bind lib. + ( + "def outer():\n" + " import numpy as lib\n" + " def middle():\n" + " [(lambda: (lib := None)) for marker in ()]\n" + " def inner():\n" + " nonlocal lib\n" + " lib.random.random()\n" + " lib = None\n" + " inner()\n" + " middle()\n", + "random source", + ), + ("if flag:\n hash = str\nkey = hash('x')\n", "builtin hash"), + ("key = hash('x')\nhash = str\n", "builtin hash"), + # A deletion, or the end of an except clause, may unbind a shadow. + ("hash = str\ndel hash\nkey = hash('x')\n", "builtin hash"), + ( + "class A:\n hash = str\n del hash\n key = hash('x')\n", + "builtin hash", + ), + ( + "hash = str\n" + "try:\n" + " pass\n" + "except ValueError as hash:\n" + " pass\n" + "key = hash('x')\n", + "builtin hash", + ), + ( + "hash = str\n" + "def clear():\n" + " global hash\n" + " del hash\n" + "clear()\n" + "key = hash('x')\n", + "builtin hash", + ), + # A function's del of a global may run after a later binding. + ( + "def clear():\n" + " global hash\n" + " del hash\n" + "hash = str\n" + "clear()\n" + "key = hash('x')\n", + "builtin hash", + ), + ( + "def clear():\n" + " global hash\n" + " if flag:\n" + " del hash\n" + "hash = str\n" + "clear()\n" + "key = hash('x')\n", + "builtin hash", + ), + ( + "def outer():\n" + " def clear():\n" + " global hash\n" + " del hash\n" + " return clear\n" + "hash = str\n" + "outer()()\n" + "def consumer(value):\n" + " return hash(value)\n", + "builtin hash", + ), + # Code may run before the statement creating it binds its names: a + # decorator may call the function, a class body runs before its + # name is bound, and an assignment's value may call a lambda in it. + ( + "def eager(f):\n" + " return f('x')\n" + "@eager\n" + "def hash(value):\n" + " return hash(value)\n", + "builtin hash", + ), + ("class hash:\n answer = hash('x')\n", "builtin hash"), + ("hash = (lambda value: hash(value))('x')\n", "builtin hash"), + # 3.14 evaluates annotations when they are read (PEP 649), after any + # later deletion; up to 3.13 they run with their statement. + ("hash = str\ndef f(x: hash('x')):\n pass\ndel hash\n", "builtin hash"), + ("hash = str\ndef f() -> hash('x'):\n pass\ndel hash\n", "builtin hash"), + ( + "class C:\n" + " hash = str\n" + " def f(self, x: hash('x')):\n" + " pass\n" + " del hash\n", + "builtin hash", + ), + ( + "class C:\n hash = str\n y: hash('x') = 1\n del hash\n", + "builtin hash", + ), + ("hash = str\ny: hash('x') = 1\ndel hash\n", "builtin hash"), + ("def hash(x: hash('x')):\n pass\n", "builtin hash"), + # A statement binds its names only after evaluating them. + ("hash = (\n hash('x')\n)\n", "builtin hash"), + ("def hash(\n value=hash('x'),\n):\n return value\n", "builtin hash"), + ("class hash(\n make(hash('x')),\n):\n pass\n", "builtin hash"), + ("(\n hash\n) = (\n hash('x')\n)\n", "builtin hash"), + # An annotation alone binds no value in a module or class body. + ("hash: int\nkey = hash('x')\n", "builtin hash"), + ( + "import numpy as lib\nclass C:\n lib: object\n answer = lib.random.uniform()\n", + "random source", + ), + # A function may run before a later binding, or after a deletion. + ( + "def consumer(value):\n" + " return hash(value)\n" + "answer = consumer('x')\n" + "hash = str\n", + "builtin hash", + ), + ("hash = str\ndef f(x):\n return hash(x)\ndel hash\n", "builtin hash"), + # A walrus in a generator expression binds around it. + ( + "import numpy as np\nlist((lib := np) for _ in 'x')\nlib.random.uniform()\n", + "random source", + ), + ("import importlib\nimportlib.import_module('random')\n", "stochastic import"), + ( + "import importlib\nimportlib.import_module(name='random')\n", + "stochastic import", + ), + ("__import__('numpy.random')\n", "stochastic import"), + ("__import__('numpy', fromlist=['random'])\n", "stochastic import"), + # Python iterates a fromlist, unpacked entries too. + ("__import__('numpy', fromlist={**{'random': None}})\n", "stochastic import"), + ("__import__('numpy', fromlist=[*('random',)])\n", "stochastic import"), + ("__import__('numpy', fromlist={*'x', *{'random': 0}})\n", "stochastic import"), + # At any depth. + ( + "__import__('numpy', fromlist={**{**{'random': None}}})\n", + "stochastic import", + ), + ("__import__('numpy', fromlist=[*[*('random',)]])\n", "stochastic import"), + ("from importlib import import_module as load\nload('zlib')\n", "hash library"), + # Hashes. + ('key = hash("x")\n', "builtin hash"), + ("keys = map(hash, names)\n", "builtin hash"), + ("import builtins\nbuiltins.hash(1)\n", "builtin hash"), + ( + "class A:\n hash = 1\n def m(self, x):\n return hash(x)\n", + "builtin hash", + ), + ("import hashlib\n", "hash library"), + ("from zlib import crc32\n", "hash library"), + ("import hashlib\ndef f(b):\n return hashlib.sha256(b)\n", "hash use"), + ("from zlib import crc32\ncrc32(b'x') % 12\n", "hash use"), + ("pd.util.hash_pandas_object(df)\n", "hash use"), + ("m = np.uint64(0x94D049BB133111EB)\n", "hash constant"), + ("gamma = 11400714819323198485\n", "hash constant"), + ("h = 0xCBF29CE484222325\n", "hash constant"), + ("x = x ^ (x >> 30)\n", "xorshift"), + ("x = (x << 13) ^ x\n", "xorshift"), + ("x ^= x >> 30\n", "xorshift"), + ("x = np.bitwise_xor(x, np.right_shift(x, 30))\n", "xorshift"), + # Sampling and ids. + ("df = df.sample(frac=0.5, random_state=1)\n", "sampling call"), + ("df = df.sample(5)\n", "sampling call"), + ("rng.shuffle(values)\n", "sampling call"), + ("pick = rng.choice(5)\n", "sampling call"), + ("month = person_id % 12\n", "id modulo"), + ('half = person("person_id", period) % 2\n', "id modulo"), + ("person_id %= 12\n", "id modulo"), + ("month = np.mod(person_id, 12)\n", "id modulo"), + ("month, _ = divmod(person_id, 12)\n", "id modulo"), + # A call of a function that draws. + ( + "def mix(x):\n return x ^ (x >> 3)\ndef f(x):\n return mix(x)\n", + "draw helper call", + ), + ], +) +def test_detector_flags_each_draw_pattern(source, kind): + assert kind in _kinds(source) + + +@pytest.mark.parametrize( + "source", + [ + "from typing import Generator\ndef f(Generator):\n return Generator\n", + "random = 3\nrandom + 1\n", + "def f(random):\n return random.random()\n", + "def hash(x):\n return x\nhash(1)\n", + # A binding sure to run before a function's def hides the builtin + # from its body; so does a binding after a deletion. + "def hash(x):\n return x\ndef f(x):\n return hash(x)\n", + "hash = str\nclass A:\n def m(self, x):\n return hash(x)\n", + "hash = str\ndel hash\nhash = repr\nkey = hash('x')\n", + "def f(x):\n hash = str\n del hash\n return x\n", + # Deleting a local leaves it unbound (UnboundLocalError); the + # builtin stays hidden. + "def f():\n hash = str\n del hash\n return hash('x')\n", + # Nothing can call an undecorated function before its def binds it. + "def hash(x):\n return hash(x - 1) if x else 0\nhash(3)\n", + "def f(hash):\n return hash(1)\n", + "class A:\n def hash(self):\n return 1\n def m(self):\n" + " return self.hash()\n", + '"""np.random.seed is banned"""\nlabel = "np.random"\n', + "obj.random_seed\nobj.sample()\n", + "months = np.ceil(weeks * 12 / 52)\nposition = (month - 1 - start) % 12\n", + "flags = x ^ y\nbits = x >> 3\nbig = 0xFFFFFFFFFFFFFFFF\n", + "rank = person.get_rank(person.household, person_id)\n", + "import numpy as np\nnp.random_like = 1\nvalue = np.where(a, b, c)\n", + "def f():\n np = object()\n return np.random\n", + "xs = [random for random in range(3)]\n", + "xp = object()\nxp.random.uniform()\n", + "getattr(np, 'where')(a, b, c)\n", + "import importlib\nimportlib.import_module('numpy')\n", + "def helper(x):\n return x + 1\ndef f(x):\n return helper(x)\n", + # A binding sure to have run hides the builtin. + "hash = str\nkey = hash('x')\n", + "class A:\n hash = str\n key = hash('x')\n", + "def f():\n hash = str\n return hash('x')\n", + # An annotation reads a binding sure to have run before its + # statement; a deferred one in a class body also sees the class. + "hash = str\ndef f(x: hash('x')):\n pass\n", + "class C:\n hash = str\n def f(self, x: hash('x')):\n pass\n", + "hash = str\ny: hash('x') = 1\n", + # A statement's bindings count for the statements after it. + "hash = str\nkey = (\n hash('x')\n)\n", + "class A:\n hash = str\n key = (\n hash('x')\n )\n", + "def hash(x):\n return (\n hash(x - 1)\n ) if x else 0\nhash(1)\n", + # An annotated value binds; in a function an annotation alone makes + # the name local (UnboundLocalError), hiding the builtin. + "hash: object = str\nkey = hash('x')\n", + "def f():\n hash: object\n return hash('x')\n", + # The walrus value is the comprehension's own np, not numpy. + "import numpy as np\n[(lib := np) for np in [object()]]\nlib.random.uniform()\n", + # nonlocal stops at the nearest function binding the name. + "def outer():\n" + " import numpy as lib\n" + " def middle():\n" + " lib = object()\n" + " def inner():\n" + " nonlocal lib\n" + " return lib.random.random()\n" + " return inner\n" + " return middle\n", + # A walrus in a lambda's default binds where the comprehension runs, + # so in middle. + "def outer():\n" + " import numpy as lib\n" + " def middle():\n" + " [(lambda x=(lib := None): x) for marker in ()]\n" + " def inner():\n" + " nonlocal lib\n" + " return lib.random.random()\n" + " return inner\n" + " return middle\n", + ], +) +def test_detector_skips_deterministic_code(source): + assert draw_findings(source) == [] + + +# Python runs these when they are used, not where they are written, so a +# later deletion of a shadowing global counts; a list comprehension runs at +# once. Each needs the syntax of the Python 3 minor version given. +@pytest.mark.parametrize( + "source, reaches_builtin, minor", + _on_this_python( + [ + ("hash = str\ng = (hash(v) for v in 'x')\ndel hash\nlist(g)\n", True, 11), + ("hash = str\ng = (hash(v) for v in 'x')\nlist(g)\n", False, 11), + ("hash = str\nkeys = [hash(v) for v in 'x']\ndel hash\n", False, 11), + ("hash = str\ntype X = hash('x')\ndel hash\nX.__value__\n", True, 12), + ("hash = str\ntype X = hash('x')\nX.__value__\n", False, 12), + ( + "hash = str\n" + "type X[T: hash('x')] = list[T]\n" + "del hash\n" + "X.__type_params__[0].__bound__\n", + True, + 12, + ), + ( + "hash = str\n" + "def f[T: hash('x')]():\n" + " pass\n" + "del hash\n" + "f.__type_params__[0].__bound__\n", + True, + 12, + ), + ( + "hash = str\n" + "class C[T: hash('x')]:\n" + " pass\n" + "del hash\n" + "C.__type_params__[0].__bound__\n", + True, + 12, + ), + ( + "class C:\n" + " hash = str\n" + " def m[T: hash('x')](self):\n" + " pass\n" + " del hash\n" + "C.m.__type_params__[0].__bound__\n", + True, + 12, + ), + ( + "hash = str\n" + "def f[T = hash('x')]():\n" + " pass\n" + "del hash\n" + "f.__type_params__[0].__default__\n", + True, + 13, + ), + ] + ), +) +def test_code_python_runs_when_used_sees_later_deletions( + source, reaches_builtin, minor +): + expected = {"builtin hash"} if reaches_builtin else set() + assert {kind for _, _, kind in draw_findings(source)} == expected + + +# Each is the body of a function outer that binds hash = str. A class body +# that binds hash anywhere, even by an annotation alone, a del, or a later, +# conditional or failed binding, reads it from its own namespace, then the +# module's, so here the builtin; so does what the body evaluates, its +# methods' defaults and annotations and its first comprehension iterables +# included. Code nested in the class reads outer's. Each needs the syntax of +# the Python 3 minor version given. +@pytest.mark.parametrize( + "body, reaches_builtin, minor", + _on_this_python( + [ + (" class C:\n key = hash('x')\n hash = int\n", True, 11), + (" class C:\n hash: int\n key = hash('x')\n", True, 11), + ( + " class C:\n" + " if False:\n" + " hash = int\n" + " key = hash('x')\n", + True, + 11, + ), + ( + " class C:\n" + " hash = int\n" + " del hash\n" + " key = hash('x')\n", + True, + 11, + ), + ( + " class C:\n" + " try:\n" + " raise ValueError\n" + " except ValueError as hash:\n" + " pass\n" + " key = hash('x')\n", + True, + 11, + ), + ( + " class C:\n" + " try:\n" + " del hash\n" + " except NameError:\n" + " pass\n" + " key = hash('x')\n", + True, + 11, + ), + ( + " class C:\n" + " try:\n" + " from no_such_module_anywhere import hash\n" + " except ImportError:\n" + " pass\n" + " key = hash('x')\n", + True, + 11, + ), + ( + " class C:\n" + " hash: int\n" + " y: hash('x')\n" + " C.__annotations__\n", + True, + 11, + ), + ( + " class C:\n" + " hash: int\n" + " def m(self, x: hash('x')):\n" + " pass\n" + " C.m.__annotations__\n", + True, + 11, + ), + ( + " class C:\n" + " hash: int\n" + " def m(self, x=hash('x')):\n" + " pass\n", + True, + 11, + ), + ( + " class C:\n" + " hash: int\n" + " keys = [v for v in [hash('x')]]\n", + True, + 11, + ), + ( + " class C:\n" + " hash: int\n" + " def m[T: hash('x')](self):\n" + " pass\n" + " C.m.__type_params__[0].__bound__\n", + True, + 12, + ), + ( + " class C:\n" + " hash: int\n" + " type X = hash('x')\n" + " C.X.__value__\n", + True, + 12, + ), + # Past a nonlocal statement around the class. + ( + " def middle():\n" + " nonlocal hash\n" + " class C:\n" + " hash: int\n" + " key = hash('x')\n" + " middle()\n", + True, + 11, + ), + # The class binds hash first, or not at all, or the read is in code + # nested in the class: outer's str. + (" class C:\n hash = str\n key = hash('x')\n", False, 11), + (" class C:\n key = hash('x')\n", False, 11), + ( + " class C:\n" + " hash: int\n" + " def m(self):\n" + " return hash('x')\n" + " C().m()\n", + False, + 11, + ), + ( + " class C:\n hash: int\n keys = [hash(v) for v in 'x']\n", + False, + 11, + ), + ] + ), +) +def test_a_class_body_that_binds_a_name_reads_it_from_the_module( + body, reaches_builtin, minor +): + findings = draw_findings("def outer():\n hash = str\n" + body) + expected = {"builtin hash"} if reaches_builtin else set() + assert {kind for _, _, kind in findings} == expected + + +def test_a_reform_factory_class_reads_a_name_it_annotates_from_the_module(): + # An annotation alone makes lib local to the class body, which then + # reads numpy from the module, not the factory's None. + source = ( + "import numpy as lib\n" + "def create_reform():\n" + " lib = None\n" + " class reform_variable:\n" + " lib: object\n" + " value = lib.random.uniform()\n" + " return reform_variable\n" + ) + assert [kind for _, _, kind in draw_findings(source)] == ["random source"] + + +# A walrus binds only if it runs: in a comprehension it may run any number +# of times, in a generator expression only as it is iterated, past the first +# operand of and, or or a chained comparison only if those reach it, and in +# a conditional expression only on its branch. So these look further out, +# here to the builtin. A walrus where it is sure to run binds. +@pytest.mark.parametrize( + "source, reaches_builtin", + [ + ("[(hash := str) for _ in ()]\nkey = hash('x')\n", True), + ("g = ((hash := str) for _ in 'x')\nkey = hash('x')\nlist(g)\n", True), + ("[v for v in '' if (hash := str)]\nkey = hash('x')\n", True), + ("{(hash := str): v for v in ()}\nkey = hash('x')\n", True), + ("flag = False\nflag and (hash := str)\nkey = hash('x')\n", True), + ("flag = True\nflag or (hash := str)\nkey = hash('x')\n", True), + ("1 < 0 < (hash := str)\nkey = hash('x')\n", True), + ("(hash := str) if False else None\nkey = hash('x')\n", True), + ("(hash := str)\nkey = hash('x')\n", False), + ("(hash := str) and None\nkey = hash('x')\n", False), + ("None if (hash := str) else None\nkey = hash('x')\n", False), + ("[] == (hash := str)\nkey = hash('x')\n", False), + ], +) +def test_a_walrus_that_may_not_run_binds_nothing_surely(source, reaches_builtin): + expected = {"builtin hash"} if reaches_builtin else set() + assert {kind for _, _, kind in draw_findings(source)} == expected + + +def test_a_walrus_in_an_assert_binds_nothing_surely(): + # python -O drops the assert. + source = "assert (hash := str)\nkey = hash('x')\n" + assert [kind for _, _, kind in draw_findings(source)] == ["builtin hash"] + + +# A type alias needs Python 3.12 syntax. +if sys.version_info >= (3, 12): + + def test_a_type_alias_binds_its_name(): + # The alias isn't callable, so Python raises before any builtin runs. + source = "type hash = str\ntry:\n hash('x')\nexcept TypeError:\n pass\n" + assert draw_findings(source) == [] + + +# Each is the body of a function outer that binds hash = str. A global +# statement in a function sends hash to the module, and so to the builtin, +# in all the code nested in that function, classes and their methods +# included. One in a class body does so in that body and what it evaluates +# (its annotations, and its methods' defaults and annotations), but not in +# the code nested in it. Python 3.11 to 3.14 read each the same way. +@pytest.mark.parametrize( + "body, reaches_builtin", + [ + ( + " def middle():\n" + " global hash\n" + " def f(x: hash('x')):\n" + " pass\n" + " f.__annotations__\n" + " middle()\n", + True, + ), + ( + " def middle():\n" + " global hash\n" + " def f() -> hash('x'):\n" + " pass\n" + " f.__annotations__\n" + " middle()\n", + True, + ), + ( + " def middle():\n" + " global hash\n" + " def f(x=hash('x')):\n" + " pass\n" + " middle()\n", + True, + ), + ( + " def middle():\n" + " global hash\n" + " def f():\n" + " return hash('x')\n" + " f()\n" + " middle()\n", + True, + ), + ( + " def middle():\n" + " global hash\n" + " (lambda: hash('x'))()\n" + " middle()\n", + True, + ), + ( + " def middle():\n" + " global hash\n" + " [hash(v) for v in 'x']\n" + " list(hash(v) for v in 'x')\n" + " middle()\n", + True, + ), + ( + " def middle():\n" + " global hash\n" + " class C:\n" + " key = hash('x')\n" + " middle()\n", + True, + ), + ( + " def middle():\n" + " global hash\n" + " class C:\n" + " y: hash('x') = 1\n" + " C.__annotations__\n" + " middle()\n", + True, + ), + ( + " def middle():\n" + " global hash\n" + " class C:\n" + " def m(self):\n" + " return hash('x')\n" + " C().m()\n" + " middle()\n", + True, + ), + ( + " def middle():\n" + " global hash\n" + " def inner():\n" + " def f():\n" + " return hash('x')\n" + " f()\n" + " inner()\n" + " middle()\n", + True, + ), + # inner's global skips middle's binding too. + ( + " def middle():\n" + " hash = str\n" + " def inner():\n" + " global hash\n" + " def f():\n" + " return hash('x')\n" + " f()\n" + " inner()\n" + " middle()\n", + True, + ), + ( + " class C:\n" + " global hash\n" + " key = hash('x')\n" + " keys = [v for v in [hash('x')]]\n", + True, + ), + ( + " class C:\n" + " global hash\n" + " y: hash('x')\n" + " C.__annotations__\n", + True, + ), + ( + " class C:\n" + " global hash\n" + " def m(self, x: hash('x')) -> hash('y'):\n" + " pass\n" + " C.m.__annotations__\n", + True, + ), + # Not in the code nested in a class body. + ( + " class C:\n" + " global hash\n" + " def m(self):\n" + " return hash('x')\n" + " C().m()\n", + False, + ), + ( + " class C:\n" + " global hash\n" + " keys = [hash(v) for v in 'x']\n", + False, + ), + ( + " class C:\n" + " global hash\n" + " class D:\n" + " key = hash('x')\n", + False, + ), + ( + " class C:\n" + " global hash\n" + " y: (lambda: hash('x'))()\n" + " C.__annotations__\n", + False, + ), + # nonlocal names outer's hash; a nested binding is the nested code's. + ( + " def middle():\n" + " nonlocal hash\n" + " def f():\n" + " return hash('x')\n" + " f()\n" + " middle()\n", + False, + ), + ( + " def middle():\n" + " global hash\n" + " def f():\n" + " hash = str\n" + " return hash('x')\n" + " f()\n" + " middle()\n", + False, + ), + ], +) +def test_global_and_nonlocal_reach_the_code_nested_in_their_scope( + body, reaches_builtin +): + findings = draw_findings("def outer():\n hash = str\n" + body) + expected = {"builtin hash"} if reaches_builtin else set() + assert {kind for _, _, kind in findings} == expected + + +def test_draw_helpers_are_followed_through_calls(): + source = ( + "def mix(x):\n" + " return x ^ (x >> 3)\n" + "def month(x):\n" + " return mix(x) % 12\n" + "def formula(x):\n" + " return month(x) + 1\n" + ) + found = {(scope, kind) for _, scope, kind in draw_findings(source)} + assert ("month", "draw helper call") in found + assert ("formula", "draw helper call") in found + # A function whose findings are exempt is not a draw helper. + exempt = {("mix", "xorshift"): 1} + found = {(scope, kind) for _, scope, kind in draw_findings(source, exempt=exempt)} + assert found == {("mix", "xorshift")} + + +def test_allowlist_counts_occurrences_per_scope(): + source = ( + "def first():\n" + " return np.random.random()\n" + "def second():\n" + " return np.random.random()\n" + ) + counted = count_findings({"m.py": draw_findings(source)}) + one = {"m.py": {("first", "random source"): 1}} + # An occurrence in a new function of an allowlisted file fails. + violations, stale = check_allowlists(counted, one, {}) + assert len(violations) == 1 and "second" in violations[0] and not stale + # A second occurrence in an allowlisted scope fails. + twice = count_findings({"m.py": draw_findings(source.replace("second", "first"))}) + violations, _ = check_allowlists(twice, one, {}) + assert violations and "2 found, 1 allowed" in violations[0] + # An entry that allows more than it matches is stale. + _, stale = check_allowlists(counted, {"m.py": {("first", "random source"): 2}}, {}) + assert stale and "2 allowed, 1 found" in stale[0] + + +def test_reusing_an_allowed_hash_or_draw_spends_an_occurrence(): + checksum = ( + "import hashlib\n" + "def checksum(data):\n" + " return hashlib.sha256(data).hexdigest()\n" + ) + exempt = {("", "hash library"): 1, ("checksum", "hash use"): 1} + + def check(source, ratchet=None): + findings = draw_findings(source, "m.py", exempt=exempt) + counted = count_findings({"m.py": findings}) + return check_allowlists(counted, {"m.py": ratchet or {}}, {"m.py": exempt}) + + assert check(checksum) == ([], []) + # A new function hashing ids through the exempt import fails. + reuse = checksum + ( + "def draw(person_id):\n" + " return hashlib.sha256(bytes(person_id)).digest()[0] % 12\n" + ) + violations, _ = check(reuse) + assert any("in draw: hash use" in violation for violation in violations) + # So does a new call of a ratcheted draw. + legacy = checksum + ( + "def legacy(x):\n" + " return x ^ (x >> 3)\n" + "def stored(x):\n" + " return legacy(x)\n" + ) + ratchet = {("legacy", "xorshift"): 1, ("stored", "draw helper call"): 1} + assert check(legacy, ratchet) == ([], []) + violations, _ = check(legacy + "def more(x):\n return legacy(x)\n", ratchet) + assert any("in more: draw helper call" in violation for violation in violations) + + +def _package_findings(sources: dict, order=None, exempt=None) -> dict: + """{module: findings} of a package ``pkg`` given {module: source}. + + ``order`` is the order the modules are analysed in. + """ + modules = Modules( + { + f"pkg.{name}": ("pkg", lambda source=source: source) + for name, source in sources.items() + } + ) + analyses = { + name: analyse(sources[name], f"{name}.py", f"pkg.{name}", "pkg", modules) + for name in order or sources + } + helpers = draw_helpers(analyses, exempt or {}) + return { + name: sorted(analysis.findings + helper_calls(analysis, helpers)) + for name, analysis in analyses.items() + } + + +def _package_kinds(sources: dict, module: str = "user") -> list: + return [kind for _, _, kind in _package_findings(sources)[module]] + + +def test_named_imports_and_reexports_resolve_to_what_they_name(): + api = "import numpy as lib\nfrom numpy import random as draws\n" + user = ( + "from .api import lib\n" + "lib.random.uniform()\n" + "from .api import draws as d\n" + "d.uniform()\n" + ) + assert _package_kinds({"api": api, "user": user}) == ["random source"] * 2 + # A draw helper re-exported by name is still a draw helper. + helper = "def mix(x):\n return x ^ (x >> 3)\n" + found = _package_findings( + { + "a": helper, + "b": "from .a import mix\n", + "c": "from .b import mix\ndef consumer(x):\n return mix(x)\n", + } + ) + assert ("consumer", "draw helper call") in {(s, k) for _, s, k in found["c"]} + + +def test_module_attributes_resolve_to_what_the_module_binds(): + # A name passed on as an attribute of another module. + sources = { + "api": "import numpy as lib\n", + "bridge": "from . import api\nlib = api.lib\n", + "user": "from .bridge import lib\nlib.random.uniform()\n", + } + assert _package_kinds(sources) == ["random source"] + # A draw helper reached as an attribute of a module that re-exports it. + found = _package_findings( + { + "a": "def mix(x):\n return x ^ (x >> 3)\n", + "b": "from .a import mix\n", + "c": "from . import b\ndef consumer(x):\n return b.mix(x)\n", + } + ) + assert ("consumer", "draw helper call") in {(s, k) for _, s, k in found["c"]} + + +def test_a_package_that_rebinds_its_submodules_name_resolves(): + # Both __init__ files bind w to the function pkg.w.w defines, so the + # attribute pkg.w is that function, though "pkg.w" also names a package. + sources = { + "pkg": ("pkg", "from .w import w\n"), + "pkg.w": ("pkg.w", "from .w import w\n"), + "pkg.w.w": ("pkg.w", "def w(x):\n return x ^ (x >> 3)\n"), + "pkg.user": ("pkg", "import pkg\ndef consumer(x):\n return pkg.w(x)\n"), + } + # This resolution needs depth 4. A low limit makes one that expands + # pkg.w.w.w into ever longer paths fail fast, with a recorded cut, + # instead of searching for minutes. + modules = Modules( + { + module: (package, lambda source=source: source) + for module, (package, source) in sources.items() + }, + max_depth=5, + ) + analyses = { + module: analyse(source, f"{module}.py", module, package, modules) + for module, (package, source) in sources.items() + } + helpers = draw_helpers(analyses, {}) + assert "pkg.w.w.w" in helpers + calls = helper_calls(analyses["pkg.user"], helpers) + assert [(scope, kind) for _, scope, kind in calls] == [ + ("consumer", "draw helper call") + ] + assert not modules.exhausted + + +def test_lookups_cut_by_the_depth_limit_are_recorded(): + # lib passes along 12 re-exports, each read where it arrives. From the + # consumer end, its lookup nests past a limit of 4 and is cut; from the + # numpy end, kept results keep each lookup shallow. A cut is recorded, so + # the package scan fails instead of depending on the order. + sources = {"m0": "import numpy as lib\n"} + sources |= {f"m{i}": f"from .m{i - 1} import lib\nlib\n" for i in range(1, 13)} + sources["user"] = "from .m12 import lib\nlib.random.uniform()\n" + consumer_first = list(reversed(sources)) + for order, max_depth, cut in [ + (consumer_first, 4, True), + (list(sources), 4, False), + (consumer_first, MAX_RESOLUTION_DEPTH, False), + ]: + modules = Modules( + { + f"pkg.{name}": ("pkg", lambda source=source: source) + for name, source in sources.items() + }, + max_depth=max_depth, + ) + analyses = { + name: analyse(sources[name], f"{name}.py", f"pkg.{name}", "pkg", modules) + for name in order + } + assert bool(modules.exhausted) == cut + if not cut: + kinds = [kind for _, _, kind in analyses["user"].findings] + assert kinds == ["random source"] + + +@pytest.mark.parametrize( + "api", + [ + "import numpy as lib\nleft, right = lib, 1\n", + "import numpy as _lib\n__all__ = ['left'] + []\nleft = _lib\n", + "import numpy as _lib\n__all__ = []\n__all__.append('left')\nleft = _lib\n", + "import numpy as _lib\n__all__ = ('x',)\n__all__ += ['left']\nleft = _lib\n", + "import numpy as _lib\n__all__ = names()\nleft = _lib\n", + ], +) +def test_star_imports_bind_what_the_module_exports(api): + user = "from .api import *\nleft.random.uniform()\n" + assert _package_kinds({"api": api, "user": user}) == ["random source"] + + +def test_star_imports_bind_private_names_only_through_all(): + user = "from .api import *\n_lib.random.uniform()\n" + assert _package_kinds({"api": "import numpy as _lib\n", "user": user}) == [] + api = "import numpy as _lib\n__all__ = ['_lib']\n" + assert _package_kinds({"api": api, "user": user}) == ["random source"] + + +def test_star_imports_see_an_earlier_all_in_a_cycle(): + # Importing api first runs user while __all__ still lists lib. + api = "import numpy as lib\n__all__ = ['lib']\nfrom . import user\n__all__ = []\n" + user = "from .api import *\nlib.random.uniform()\n" + assert _package_kinds({"api": api, "user": user}) == ["random source"] + # Or before api sets __all__, when every public name it has bound is + # exported. + api = "import numpy as lib\nfrom . import user\n__all__ = []\n" + assert _package_kinds({"api": api, "user": user}) == ["random source"] + # A private name is exported only while an earlier __all__ lists it. + api = "import numpy as _lib\n__all__ = ['_lib']\nfrom . import user\n__all__ = []\n" + user = "from .api import *\n_lib.random.uniform()\n" + assert _package_kinds({"api": api, "user": user}) == ["random source"] + + +@pytest.mark.parametrize( + "api", + [ + # Importing api first runs user before __all__ lists hash, or before + # api binds it. + "__all__ = []\nfrom . import user\nhash = str\n__all__ = ['hash']\n", + "hash = str\n__all__ = []\nfrom . import user\n__all__ = ['hash']\n", + "from . import user\nhash = str\n", + # Through another module, a package's __init__, a function that + # imports, or a literal dynamic import. + "from . import other\nhash = str\n", + "import pkg.sub.leaf\nhash = str\n", + "def load():\n from . import user\nload()\nhash = str\n", + "import importlib\nimportlib.import_module('pkg.user')\nhash = str\n", + ], +) +def test_star_imports_in_an_import_cycle_bind_nothing_surely(api): + sources = { + "api": api, + "other": "from . import user\n", + "sub": "from pkg import user\n", + "sub.leaf": "", + "user": "from .api import *\nkey = hash('x')\n", + } + assert _package_kinds(sources) == ["builtin hash"] + # A module's star import of itself sees itself part-run too. + itself = {"user": "from .user import *\nkey = hash('x')\nhash = str\n"} + assert _package_kinds(itself) == ["builtin hash"] + + +def _modules_of(sources: dict) -> Modules: + """Modules for {module: source}; a module with submodules is a package.""" + packages = { + module: module + if any(other.startswith(module + ".") for other in sources) + else module.rsplit(".", 1)[0] + for module in sources + } + return Modules( + { + module: (packages[module], lambda source=source: source) + for module, source in sources.items() + } + ) + + +def _assert_users_find(sources: dict, reaches_builtin: bool): + """Each module that reads ``hash('x')`` finds the builtin, or not, in + every order the modules are analysed in.""" + users = [module for module, source in sources.items() if "hash('x')" in source] + expected = ["builtin hash"] if reaches_builtin else [] + for order in permutations(sources): + modules = _modules_of(sources) + analyses = { + module: analyse( + sources[module], + f"{module}.py", + module, + modules.sources[module][0], + modules, + ) + for module in order + } + for user in users: + assert [kind for _, _, kind in analyses[user].findings] == expected + + +# Some import order starts user while api is part-run: through the modules +# a star import of a package loads (what its __all__ lists, in any version +# of it, dotted names too, with the packages on the way, or every module +# below it where __all__ is built otherwise), or a dynamic import under +# another name or relative to the package its arguments give. In the +# controls no order does: without __all__ a star import of a package loads +# none of its submodules, one of the package around the importer starts +# nothing, as that has run, a dotted name loads only what it names, and a +# relative import's own package is the one its globals give. +@pytest.mark.parametrize( + "sources, reaches_builtin", + [ + ( + { + "pkg": "__all__ = ['user']\n", + "pkg.api": "from . import *\nhash = str\n", + "pkg.user": "from .api import *\nkey = hash('x')\n", + }, + True, + ), + ( + { + "pkg.sub": "__all__ = ['leaf']\n", + "pkg.sub.leaf": "from .. import user\n", + "pkg.api": "from .sub import *\nhash = str\n", + "pkg.user": "from .api import *\nkey = hash('x')\n", + }, + True, + ), + ( + { + "pkg.sub": "__all__ = ['user']\n", + "pkg.sub.user": "from pkg.api import *\nkey = hash('x')\n", + "pkg.api": "from pkg.sub import *\nhash = str\n", + }, + True, + ), + # Imported first, pkg.sub imports api while its __all__ lists leaf. + ( + { + "pkg.sub": "__all__ = ['leaf']\nfrom .. import api\n__all__ = []\n", + "pkg.sub.leaf": "from .. import user\n", + "pkg.api": "from .sub import *\nhash = str\n", + "pkg.user": "from .api import *\nkey = hash('x')\n", + }, + True, + ), + ( + { + "pkg.sub": "names = ['leaf']\n__all__ = list(names)\n", + "pkg.sub.leaf": "from .. import user\n", + "pkg.api": "from .sub import *\nhash = str\n", + "pkg.user": "from .api import *\nkey = hash('x')\n", + }, + True, + ), + ( + { + "pkg.api": "from importlib import import_module as load\n" + "load('pkg.user')\n" + "hash = str\n", + "pkg.user": "from .api import *\nkey = hash('x')\n", + }, + True, + ), + ( + { + "pkg.api": "import importlib\n" + "getattr(importlib, 'import_module')('pkg.user')\n" + "hash = str\n", + "pkg.user": "from .api import *\nkey = hash('x')\n", + }, + True, + ), + ( + { + "pkg.sub": "", + "pkg.sub.leaf": "from .. import user\n", + "pkg.api": "from .sub import *\nhash = str\n", + "pkg.user": "from .api import *\nkey = hash('x')\n", + }, + False, + ), + # pkg.sub started before api, so api's star import of pkg starts no + # module again. Python cannot bind sub, part-run, and api catches it. + ( + { + "pkg": "__all__ = ['sub']\n", + "pkg.sub": "from . import user\n", + "pkg.sub.api": "try:\n" + " from .. import *\n" + "except AttributeError:\n" + " pass\n" + "hash = str\n", + "pkg.sub.user": "from .api import *\nkey = hash('x')\n", + }, + False, + ), + ( + { + "pkg.sub": "__all__ = ['user']\n", + "pkg.sub.user": "from pkg.api import *\nkey = hash('x')\n", + "pkg.api": "__import__('pkg.sub', fromlist=['*'])\nhash = str\n", + }, + True, + ), + ( + { + "pkg.sub": "__all__ = ['user']\n", + "pkg.sub.user": "from pkg.api import *\nkey = hash('x')\n", + "pkg.api": "from builtins import __import__ as load\n" + "load('pkg.sub', fromlist=('*',))\n" + "hash = str\n", + }, + True, + ), + ( + { + "pkg.sub": "__all__ = ['user']\n", + "pkg.sub.user": "from pkg.api import *\nkey = hash('x')\n", + "pkg.api": "__import__('sub', globals(), None, ['*'], 1)\nhash = str\n", + }, + True, + ), + ( + { + "pkg.sub": "__all__ = ['user']\n", + "pkg.sub.user": "from pkg.api import *\nkey = hash('x')\n", + "pkg.api": "__import__('pkg.sub', fromlist='*')\nhash = str\n", + }, + True, + ), + ( + { + "pkg.sub": "", + "pkg.sub.user": "from pkg.api import *\nkey = hash('x')\n", + "pkg.api": "__import__('pkg.sub', fromlist=['*'])\nhash = str\n", + }, + False, + ), + ( + { + "pkg.sub": "__all__ = ['branch.user']\n", + "pkg.sub.branch": "", + "pkg.sub.branch.user": "from pkg.api import *\nkey = hash('x')\n", + "pkg.api": "__import__('pkg.sub', fromlist=['*'])\nhash = str\n", + }, + True, + ), + # Python loads pkg.sub.branch.user, then fails to bind 'branch.user'. + ( + { + "pkg.sub": "__all__ = ['branch.user']\n", + "pkg.sub.branch": "", + "pkg.sub.branch.user": "from pkg.api import *\nkey = hash('x')\n", + "pkg.api": "try:\n" + " from .sub import *\n" + "except AttributeError:\n" + " pass\n" + "hash = str\n", + }, + True, + ), + ( + { + "pkg.sub": "__all__ = ['branch.leaf']\n", + "pkg.sub.branch": "from . import user\n", + "pkg.sub.branch.leaf": "", + "pkg.sub.branch.user": "from pkg.api import *\nkey = hash('x')\n", + "pkg.api": "__import__('pkg.sub', fromlist=['*'])\nhash = str\n", + }, + True, + ), + ( + { + "pkg.sub": "names = ['branch.user']\n__all__ = list(names)\n", + "pkg.sub.branch": "", + "pkg.sub.branch.user": "from pkg.api import *\nkey = hash('x')\n", + "pkg.api": "__import__('pkg.sub', fromlist=['*'])\nhash = str\n", + }, + True, + ), + ( + { + "pkg.sub": "__all__ = ['branch.other']\n", + "pkg.sub.branch": "", + "pkg.sub.branch.other": "", + "pkg.sub.branch.user": "from pkg.api import *\nkey = hash('x')\n", + "pkg.api": "__import__('pkg.sub', fromlist=['*'])\nhash = str\n", + }, + False, + ), + ( + { + "pkg.other": "", + "pkg.other.sub": "__all__ = ['user']\n", + "pkg.other.sub.user": "from pkg.api import *\nkey = hash('x')\n", + "pkg.api": "__import__('sub', {'__package__': 'pkg.other'}, None, ['*'], 1)\n" + "hash = str\n", + }, + True, + ), + ( + { + "pkg.other": "", + "pkg.other.sub": "__all__ = ['user']\n", + "pkg.other.sub.user": "from pkg.api import *\nkey = hash('x')\n", + "pkg.api": "__import__('sub', {'__name__': 'pkg.other.mod'}, None, ['*'], 1)\n" + "hash = str\n", + }, + True, + ), + ( + { + "pkg.other": "", + "pkg.other.sub": "__all__ = ['user']\n", + "pkg.other.sub.user": "from pkg.api import *\nkey = hash('x')\n", + "pkg.api": "__import__(\n" + " 'sub', {'__name__': 'pkg.other', '__path__': []}, None, ['*'], 1\n" + ")\n" + "hash = str\n", + }, + True, + ), + # The last of a repeated key wins: __package__ is None, so __name__ + # gives the package. + ( + { + "pkg.other": "", + "pkg.other.sub": "__all__ = ['user']\n", + "pkg.other.sub.user": "from pkg.api import *\nkey = hash('x')\n", + "pkg.api": "__import__(\n" + " 'sub',\n" + " {'__package__': 'pkg', '__package__': None, '__name__': 'pkg.other.mod'},\n" + " None,\n" + " ['*'],\n" + " 1,\n" + ")\n" + "hash = str\n", + }, + True, + ), + # A dict fromlist is iterated by key. + ( + { + "pkg.other": "", + "pkg.other.sub": "__all__ = ['user']\n", + "pkg.other.sub.user": "from pkg.api import *\nkey = hash('x')\n", + "pkg.api": "__import__('pkg.other.sub', fromlist={'*': None})\nhash = str\n", + }, + True, + ), + # And so are a literal ``**`` or ``*`` entry's own items. + ( + { + "pkg.other": "", + "pkg.other.sub": "__all__ = ['user']\n", + "pkg.other.sub.user": "from pkg.api import *\nkey = hash('x')\n", + "pkg.api": "__import__('pkg.other.sub', fromlist={**{'*': None}})\n" + "hash = str\n", + }, + True, + ), + ( + { + "pkg.other": "", + "pkg.other.sub": "__all__ = ['user']\n", + "pkg.other.sub.user": "from pkg.api import *\nkey = hash('x')\n", + "pkg.api": "__import__('pkg.other.sub', fromlist=[*['*']])\nhash = str\n", + }, + True, + ), + ( + { + "pkg.other": "", + "pkg.other.sub": "__all__ = ['user']\n", + "pkg.other.sub.user": "from pkg.api import *\nkey = hash('x')\n", + "pkg.api": "__import__(\n" + " 'sub', {**{'__package__': 'pkg.other'}}, None, ['*'], 1\n" + ")\n" + "hash = str\n", + }, + True, + ), + # An empty one names nothing: no star load, so user never runs. + ( + { + "pkg.other": "", + "pkg.other.sub": "__all__ = ['user']\n", + "pkg.other.sub.user": "from pkg.api import *\nkey = hash('x')\n", + "pkg.api": "__import__('pkg.other.sub', fromlist={**{}})\nhash = str\n", + }, + False, + ), + ( + { + "pkg.api": "import importlib\n" + "importlib.import_module('.user', __package__)\n" + "hash = str\n", + "pkg.user": "from .api import *\nkey = hash('x')\n", + }, + True, + ), + ( + { + "pkg.sub": "__all__ = ['user']\n", + "pkg.sub.user": "from pkg.api import *\nkey = hash('x')\n", + "pkg.other": "", + "pkg.other.sub": "__all__ = []\n", + "pkg.api": "__import__('sub', {'__package__': 'pkg.other'}, None, ['*'], 1)\n" + "hash = str\n", + }, + False, + ), + ], +) +def test_star_imports_see_cycles_through_star_loads_and_renamed_imports( + sources, reaches_builtin +): + _assert_users_find({"pkg": ""} | sources, reaches_builtin) + + +@pytest.mark.parametrize( + "sources, reaches_builtin", + [ + ( + { + "pkg": "def load():\n from . import user\n", + "pkg.api": "from pkg import load\nload()\nhash = str\n", + "pkg.user": "from pkg.api import *\nkey = hash('x')\n", + }, + True, + ), + ( + { + "pkg": "from .helpers import go\ndef load():\n go()\n", + "pkg.helpers": "def go():\n from . import user\n", + "pkg.api": "from pkg import load\nload()\nhash = str\n", + "pkg.user": "from pkg.api import *\nkey = hash('x')\n", + }, + True, + ), + ( + { + "pkg": "def load():\n from .mid import user\n", + "pkg.mid": "", + "pkg.mid.api": "from pkg import load\nload()\nhash = str\n", + "pkg.mid.user": "from pkg.mid.api import *\nkey = hash('x')\n", + }, + True, + ), + ( + { + "pkg": "def load():\n __import__('pkg.sub', fromlist=['*'])\n", + "pkg.api": "from pkg import load\nload()\nhash = str\n", + "pkg.sub": "__all__ = ['user']\n", + "pkg.sub.user": "from pkg.api import *\nkey = hash('x')\n", + }, + True, + ), + ( + { + "pkg": "load = lambda: __import__('pkg.user')\n", + "pkg.api": "from pkg import load\nload()\nhash = str\n", + "pkg.user": "from pkg.api import *\nkey = hash('x')\n", + }, + True, + ), + ( + { + "pkg": "loads = (__import__('pkg.user') for _ in [0])\n", + "pkg.api": "from pkg import loads\nnext(loads)\nhash = str\n", + "pkg.user": "from pkg.api import *\nkey = hash('x')\n", + }, + True, + ), + ( + { + "pkg": "class Loader:\n" + " def load(self):\n" + " from . import user\n", + "pkg.api": "from pkg import Loader\nLoader().load()\nhash = str\n", + "pkg.user": "from pkg.api import *\nkey = hash('x')\n", + }, + True, + ), + ( + { + "pkg": "from .sub import *\n", + "pkg.sub": "__all__ = ['helpers']\n", + "pkg.sub.helpers": "def go():\n from .. import user\n", + "pkg.api": "from pkg import helpers\nhelpers.go()\nhash = str\n", + "pkg.user": "from pkg.api import *\nkey = hash('x')\n", + }, + True, + ), + # pkg ran its import of other before api started; nothing api may + # call imports user. + ( + { + "pkg": "from . import other\ndef load():\n return other\n", + "pkg.other": "VALUE = 1\n", + "pkg.api": "from pkg import load\nload()\nhash = str\n", + "pkg.user": "from pkg.api import *\nkey = hash('x')\n", + }, + False, + ), + # pkg.mid ran, with its import of user, before api started: a + # function's import of it starts nothing. + ( + { + "pkg": "def load():\n from . import mid\n", + "pkg.mid": "from . import user\n", + "pkg.mid.api": "from .. import load\nhash = str\n", + "pkg.mid.user": "from pkg.mid.api import *\nkey = hash('x')\n", + }, + False, + ), + ], +) +def test_star_imports_see_cycles_through_code_the_packages_around_may_run( + sources, reaches_builtin +): + # The packages around a module, and what they imported as they ran, have + # run before it, but its code may call their functions: an import in + # one is an edge. + _assert_users_find(sources, reaches_builtin) + + +# main loads pkg.helper, then api reaches it without an import, as an +# attribute of its package, and calls code in it that starts user while api +# is part-run. +@pytest.mark.parametrize( + "api, helper", + [ + ("import pkg\npkg.helper.go()\n", None), + ( + "import pkg\npkg.helper.trigger\n", + "def __getattr__(name):\n from . import user\n", + ), + ("import pkg as p\np.helper.go()\n", None), + ("import pkg\ngetattr(pkg, 'helper').go()\n", None), + ("import sys\nsys.modules['pkg.helper'].go()\n", None), + ("import sys\nsys.modules.get('pkg.helper').go()\n", None), + ("__import__('pkg').helper.go()\n", None), + ("from pkg import *\nhelper.go()\n", None), + ], +) +def test_star_imports_see_cycles_through_modules_reached_without_an_import(api, helper): + sources = { + "pkg": "", + "pkg.helper": helper or "def go():\n from . import user\n", + "pkg.main": "import pkg.helper\nimport pkg.api\n", + "pkg.api": api + "hash = str\n", + "pkg.user": "from pkg.api import *\nkey = hash('x')\n", + } + _assert_users_find(sources, True) + + +@pytest.mark.parametrize( + "sources, reaches_builtin", + [ + # An alias, assigned, of a package whose submodule main loaded. + ( + { + "pkg.sub": "", + "pkg.sub.helper": "def go():\n from .. import user\n", + "pkg.main": "import pkg.sub.helper\nimport pkg.api\n", + "pkg.api": "import pkg\nh = pkg.sub\nh.helper.go()\nhash = str\n", + "pkg.user": "from pkg.api import *\nkey = hash('x')\n", + }, + True, + ), + # An alias bound to a path below itself. + ( + { + "pkg.sub": "", + "pkg.sub.helper": "def go():\n from .. import user\n", + "pkg.main": "import pkg.sub.helper\nimport pkg.api\n", + "pkg.api": "import pkg\nh = pkg\nh = h.sub\nh.helper.go()\nhash = str\n", + "pkg.user": "from pkg.api import *\nkey = hash('x')\n", + }, + True, + ), + # A function in a module api imports reaches it. + ( + { + "pkg.helper": "def go():\n from . import user\n", + "pkg.x": "import pkg\ndef run():\n pkg.helper.go()\n", + "pkg.main": "import pkg.helper\nimport pkg.api\n", + "pkg.api": "from pkg.x import run\nrun()\nhash = str\n", + "pkg.user": "from pkg.api import *\nkey = hash('x')\n", + }, + True, + ), + # So does a function in the package around api, by its own name for + # its submodule. + ( + { + "pkg": "def load():\n helper.go()\n", + "pkg.helper": "def go():\n from . import user\n", + "pkg.main": "import pkg.helper\nimport pkg.api\n", + "pkg.api": "from pkg import load\nload()\nhash = str\n", + "pkg.user": "from pkg.api import *\nkey = hash('x')\n", + }, + True, + ), + # The code api reaches starts nothing, or nothing that starts user. + ( + { + "pkg": "from . import helper\n", + "pkg.helper": "def go():\n return 1\n", + "pkg.api": "import pkg\npkg.helper.go()\nhash = str\n", + "pkg.user": "from pkg.api import *\nkey = hash('x')\n", + }, + False, + ), + ( + { + "pkg": "from . import helper\n", + "pkg.helper": "def go():\n from . import other\n", + "pkg.other": "VALUE = 1\n", + "pkg.api": "import pkg\npkg.helper.go()\nhash = str\n", + "pkg.user": "from pkg.api import *\nkey = hash('x')\n", + }, + False, + ), + # A plain module's bare name is its own, not its package's submodule. + ( + { + "pkg.helper": "def go():\n from . import user\n", + "pkg.main": "import pkg.helper\nimport pkg.api\n", + "pkg.api": "def helper():\n pass\nhelper()\nhash = str\n", + "pkg.user": "from pkg.api import *\nkey = hash('x')\n", + }, + False, + ), + ], +) +def test_star_imports_see_cycles_through_modules_reached_by_aliases_and_functions( + sources, reaches_builtin +): + _assert_users_find({"pkg": ""} | sources, reaches_builtin) + + +def test_a_computed_all_may_name_any_module_below_its_package(): + # __all__ holds nothing when this one runs, but a computed __all__ may + # hold a dotted name, so a star import of the package may load any module + # below it, and the scan errs toward reporting. + sources = { + "pkg": "__all__ = list([])\n", + "pkg.sub": "from . import user\n", + "pkg.sub.api": "from .. import *\nhash = str\n", + "pkg.sub.user": "from .api import *\nkey = hash('x')\n", + } + _assert_users_find(sources, True) + + +@pytest.mark.parametrize( + "source, run, later", + [ + ("import pkg\npkg.a.f()\n", {"pkg.a"}, set()), + ("import pkg\ndef f():\n pkg.a.g()\n", set(), {"pkg.a"}), + ("import pkg as p\nx = p.a\nx.b.f()\n", {"pkg.a", "pkg.a.b"}, set()), + ("from pkg import a as b\nb.c.f()\n", {"pkg.a", "pkg.a.c"}, set()), + ("import pkg\ng = getattr(pkg, 'a')\n", {"pkg.a"}, set()), + ("import sys\nsys.modules['pkg.a'].f()\n", {"pkg.a"}, set()), + ("from sys import modules\nmodules.get('pkg.a').f()\n", {"pkg.a"}, set()), + ("f = lambda: __import__('pkg').a.f()\n", set(), {"pkg.a"}), + # __import__ returns the top package; import_module the module. + ("__import__('pkg.sub').a.f()\n", {"pkg.a", "pkg.sub.a"}, set()), + ("from pkg.sub import *\na.f()\n", {"pkg.sub.a"}, set()), + ("pkg.a.f()\n", {"pkg.a"}, set()), + # A name bound to a path below itself. + ("import pkg\nx = pkg\nx = x.a\nx.b.f()\n", {"pkg.a", "pkg.a.b"}, set()), + ("import pkg\nx = pkg.a\ny = x.b\nx = y\nx.f()\n", {"pkg.a", "pkg.a.b"}, set()), + # y's binding reads x's, bound after it. + ( + "import pkg\nfor i in range(2):\n y = x if i else pkg\n x = y.a\n" + "y.b.f()\n", + {"pkg.a", "pkg.a.b"}, + set(), + ), + # Either branch, inline or bound. + ( + "import pkg\nh = pkg.a if c else None\nh.b.f()\n", + {"pkg.a", "pkg.a.b"}, + set(), + ), + ("import pkg\nh = c or pkg.a\nh.c.f()\n", {"pkg.a", "pkg.a.c"}, set()), + ("import pkg\nh: object = pkg.a\nh.b.f()\n", {"pkg.a", "pkg.a.b"}, set()), + ( + "import pkg\n(h := pkg.a).b.f()\nh.c\n", + {"pkg.a", "pkg.a.b", "pkg.a.c"}, + set(), + ), + # A plain module's bare name is its own, not its package's submodule. + ("a.f()\n", set(), set()), + ], +) +def test_references_are_split_by_when_they_may_run(source, run, later): + tree = ast.parse(source) + (_, run_stars), (_, later_stars) = _imports_in(tree, "pkg", "pkg.m") + modules = {"pkg.a", "pkg.a.b", "pkg.a.c", "pkg.sub.a"} + known = {path for module in modules for path in _started(module)} + found = _references_in(tree, "pkg", run_stars | later_stars, known, "pkg.m") + assert (found[0] & modules, found[1] & modules) == (run, later) + + +@pytest.mark.parametrize( + "source", + [ + "import pkg\npkg.a.f()\nx.y.z\n", + "import pkg as p\nq = p.a\nq.b.c.f()\nimport os\nos.path.join('a')\n", + "import sys\nsys.modules['pkg.a'].f()\nsys.modules['os'].sep\ngetattr(o, 'n')\n", + ], +) +def test_reference_paths_stay_within_the_known_paths(source): + # Only paths a package module may yet be named from, and the names the + # rules read, are kept, so the binding fixpoint has finitely many paths + # to add and ends. + modules = {"pkg.a", "pkg.a.b"} + known = {path for module in modules for path in _started(module)} + run, later = _references_in(ast.parse(source), "pkg", set(), known, "pkg.m") + assert run | later <= known | REFERENCE_NAMES + + +# What Python imports for each call, read from its arguments. +@pytest.mark.parametrize( + "call, imported", + [ + ("__import__('sub', globals(), None, ['x'], 1)", ["pkg.sub", "pkg.sub.x"]), + ( + "__import__('sub', {'__package__': 'pkg.other'}, None, (), 1)", + ["pkg.other.sub"], + ), + ( + "__import__('sub', {'__package__': None, '__name__': 'pkg.o.m'}, level=1)", + ["pkg.o.sub"], + ), + ( + "__import__('sub', {'__name__': 'pkg.o', '__path__': None}, level=1)", + ["pkg.o.sub"], + ), + ("__import__('sub', {'__name__': 'pkg.o.m'}, level=2)", ["pkg.sub"]), + # The last of a repeated key wins, and None then reads as absent. + ( + "__import__('sub', {'__package__': 'pkg', '__package__': None, " + "'__name__': 'pkg.o.m'}, level=1)", + ["pkg.o.sub"], + ), + ( + "__import__('sub', {'__package__': None, '__package__': 'pkg.o'}, level=1)", + ["pkg.o.sub"], + ), + ("__import__('sub', {'__name__': 'pkg.o.m', '__name__': None}, level=1)", []), + ("__import__('sub', globals(), None, {'x': 0}, 1)", ["pkg.sub", "pkg.sub.x"]), + # Unpacked literal entries are read in place; computed ones give none. + ( + "__import__('sub', globals(), None, {**{'x': 0}}, 1)", + ["pkg.sub", "pkg.sub.x"], + ), + ( + "__import__('sub', globals(), None, [*computed, 'x'], 1)", + ["pkg.sub", "pkg.sub.x"], + ), + ("__import__('sub', {**{'__package__': 'pkg.o'}}, level=1)", ["pkg.o.sub"]), + ( + "__import__('sub', {'__package__': 'pkg', **{'__package__': None}, " + "'__name__': 'pkg.o.m'}, level=1)", + ["pkg.o.sub"], + ), + ( + "__import__('sub', {**{'__package__': 'pkg.o'}, **other}, level=1)", + ["pkg.sub"], + ), + ("__import__('sub', {1: 0, '__package__': 'pkg.o'}, level=1)", ["pkg.o.sub"]), + # At any depth. + ( + "__import__('sub', globals(), None, {**{**{'x': 0}}}, 1)", + ["pkg.sub", "pkg.sub.x"], + ), + ( + "__import__('sub', globals(), None, [*[*['x']]], 1)", + ["pkg.sub", "pkg.sub.x"], + ), + ("__import__('sub', {**{**{'__package__': 'pkg.o'}}}, level=1)", ["pkg.o.sub"]), + ( + "__import__('sub', {'__spec__': spec, '__name__': 'x.y'}, level=1)", + ["pkg.sub"], + ), + ("__import__('sub', {'__package__': computed}, level=1)", ["pkg.sub"]), + ("__import__('sub', {'__name__': name}, level=1)", ["pkg.sub"]), + ("__import__('sub', {**other}, level=1)", ["pkg.sub"]), + ("__import__('sub', mapping, level=1)", ["pkg.sub"]), + ("__import__('sub', {}, level=1)", []), + ("import_module('.sub', __package__)", ["pkg.sub"]), + ("import_module('.sub', package=__spec__.parent)", ["pkg.sub"]), + ("import_module('.sub', __name__)", ["pkg.m.sub"]), + ("import_module('..sub', 'pkg.o')", ["pkg.sub"]), + ("import_module('.sub', computed)", ["pkg.sub"]), + ("import_module('.sub')", []), + ("import_module('.sub', None)", []), + ], +) +def test_relative_dynamic_imports_resolve_as_python_reads_them(call, imported): + node = ast.parse(call, mode="eval").body + names = _dynamic_imports(node, DYNAMIC_IMPORTS, "pkg", "pkg.m") + assert [name for name in names if name.startswith("pkg.")] == imported + + +@pytest.mark.parametrize( + "source, run, later", + [ + ("import pkg.a\nclass C:\n import pkg.b\n", {"pkg.a", "pkg.b"}, set()), + ( + "if x:\n import pkg.a\ntry:\n import pkg.b\nexcept E:\n pass\n", + {"pkg.a", "pkg.b"}, + set(), + ), + ("def f():\n import pkg.a\n", set(), {"pkg.a"}), + ("async def f():\n import pkg.a\n", set(), {"pkg.a"}), + ("def f():\n if x:\n import pkg.a\n", set(), {"pkg.a"}), + ("class C:\n def m(self):\n import pkg.a\n", set(), {"pkg.a"}), + ("f = lambda: __import__('pkg.a')\n", set(), {"pkg.a"}), + ("g = (__import__('pkg.a') for _ in [0])\n", set(), {"pkg.a"}), + ("x: __import__('pkg.a')\n", set(), {"pkg.a"}), + ("def f(x: __import__('pkg.a')): pass\n", set(), {"pkg.a"}), + ("@__import__('pkg.a')\ndef f(): pass\n", set(), {"pkg.a"}), + ], +) +def test_imports_are_split_by_when_they_may_run(source, run, later): + (run_names, _), (later_names, _) = _imports_in(ast.parse(source), "pkg") + assert (run_names, later_names) == (run, later) + + +# Each needs the syntax of the Python 3 minor version given, and none parses +# before 3.12. +if sys.version_info >= (3, 12): + + @pytest.mark.parametrize( + "minor, source", + _on_this_python( + [ + (12, "class C[T: __import__('pkg.a')]:\n pass\n"), + (12, "type Alias = __import__('pkg.a')\n"), + (13, "def f[T = __import__('pkg.a')](): pass\n"), + ], + minor=0, + ), + ) + def test_type_parameters_and_aliases_may_run_later(minor, source): + (run_names, _), (later_names, _) = _imports_in(ast.parse(source), "pkg") + assert (run_names, later_names) == (set(), {"pkg.a"}) + + +def test_imports_in_name_star_bases_by_when_they_may_run(): + source = ( + "from .sub import *\n" + "__import__('pkg.top', fromlist=['*'])\n" + "def f():\n __import__('pkg.inner', fromlist='*')\n" + ) + (_, run_stars), (_, later_stars) = _imports_in(ast.parse(source), "pkg") + assert (run_stars, later_stars) == ({"pkg.sub", "pkg.top"}, {"pkg.inner"}) + + +def test_star_imports_see_all_of_a_module_whose_package_imports_the_importer(): + # Importing pkg.api starts pkg first, which runs user, whose star import + # then runs all of api: api's import of pkg starts nothing. + sources = { + "pkg": ("pkg", "from . import user\n"), + "pkg.api": ("pkg", "import pkg\nhash = str\n"), + "pkg.user": ("pkg", "from .api import *\nkey = hash('x')\n"), + } + modules = Modules( + { + module: (package, lambda source=source: source) + for module, (package, source) in sources.items() + } + ) + for module, (package, source) in sources.items(): + assert analyse(source, f"{module}.py", module, package, modules).findings == [] + + +@pytest.mark.parametrize( + "api", + [ + # An earlier __all__, which a module imported in a cycle sees. + "hash = str\n__all__ = ['hash']\nfrom . import user\n__all__ = []\n", + "hash = str\n__all__ = ['hash']\n__all__.remove('hash')\n", + "hash = str\n__all__ = ['hash']\n__all__.remove(name)\n", + "if flag:\n hash = str\n", + "hash = str\ndel hash\n", + "hash = str\n__all__ = names()\n", + ], +) +def test_star_imports_may_not_bind_some_names(api): + # So the importer looks further out, here to the builtin. + user = "from .api import *\nkey = hash('x')\n" + assert _package_kinds({"api": api, "user": user}) == ["builtin hash"] + + +@pytest.mark.parametrize( + "api", + [ + "hash = str\n", + "hash = str\n__all__ = ['hash']\n", + "from .other import hash\n", + # other doesn't import user. + "from . import other\nhash = str\n", + ], +) +def test_star_imports_bind_what_the_module_surely_exports(api): + user = "from .api import *\nkey = hash('x')\n" + sources = {"api": api, "other": "hash = str\n", "user": user} + assert _package_kinds(sources) == [] + + +def test_cyclic_star_imports_resolve_in_any_order(): + sources = { + "a": "from .b import *\nimport numpy as lib\n", + "b": "from .a import *\n", + "user": "from .b import *\nlib.random.uniform()\n", + } + results = [ + _package_findings(sources, order) + for order in (["a", "b", "user"], ["user", "b", "a"], ["b", "user", "a"]) + ] + assert results[0] == results[1] == results[2] + assert [kind for _, _, kind in results[0]["user"]] == ["random source"] + + +@pytest.mark.parametrize( + "sources", + [ + { + "a": "if flag:\n import numpy as lib\nelse:\n from .b import lib\n", + "b": "from .a import lib\n", + "other": "from .a import lib\nlib.random.uniform()\n", + "user": "from .b import lib\nlib.random.uniform()\n", + }, + # From via_a, a.lib leads to b.lib, then a.lib again, then b.lib, + # which is cut: so that inner a.lib has none of b's numpy. Keeping + # that partial answer would hide numpy from via_c, which reaches it + # only through a.lib, whenever via_a is analysed first. + { + "a": "from .b import lib\n", + "b": "if flag:\n from .a import lib\nelse:\n import numpy as lib\n", + "c": "from .a import lib\n", + "via_a": "from .a import lib\nlib.random.uniform()\n", + "via_c": "from .c import lib\nlib.random.uniform()\n", + }, + ], +) +def test_cyclic_named_imports_resolve_in_any_order(sources): + results = [ + _package_findings(sources, list(order)) for order in permutations(sources) + ] + assert all(result == results[0] for result in results) + for module, source in sources.items(): + if "lib.random" in source: + assert [kind for _, _, kind in results[0][module]] == ["random source"] + + +def test_model_api_passes_on_cores_numpy_and_its_own_helpers(): + modules = package_scan().modules + user = "policyengine_uk.variables.gov.dwp.tax_credits_current_year_income" + assert "numpy" in modules.star_binds("policyengine_uk.model_api", "np", user) + assert modules.star_binds_surely("policyengine_uk.model_api", "np", user) + for api in ("policyengine_uk.model_api", "policyengine_core.model_api"): + assert "numpy.random.uniform" in modules.expand(f"{api}.np.random.uniform") + assert "policyengine_uk.utils.benefit_unit.add_for_members" in modules.star_binds( + "policyengine_uk.model_api", "add_for_members", user + ) + + +def test_splitmix64_hash_and_its_uprating_caller_are_draw_helpers(): + helpers = package_scan().helpers + assert "policyengine_uk.utils.stochastic.splitmix64_uniform" in helpers + assert ( + "policyengine_uk.data.economic_assumptions.uprate_student_loan_plans" in helpers + ) + assert "policyengine_uk.build_metadata.get_data_build_fingerprint" not in helpers + + +def test_package_draws_nothing_outside_the_allowlists(): + violations, _ = check_allowlists(package_scan().counts, RATCHET, EXEMPT) + assert not violations, ( + "The model must not draw, sample, shuffle or hash entities to a " + "pseudo-random value while a simulation is built or run " + "(policyengine-core#500). Draw in the dataset build and store the " + "draw as an input variable; the formula then compares it with the " + "rule, so a rate stays reformable. Draw patterns found:\n" + + "\n".join(violations) + ) + + +def test_allowlist_entries_still_match(): + _, stale = check_allowlists(package_scan().counts, RATCHET, EXEMPT) + assert not stale, ( + "These entries allow more than the package now has. Lower or delete " + "them, and lower RATCHET_CEILING to match RATCHET:\n" + "\n".join(stale) + ) + + +def test_package_lookups_stay_within_the_depth_limit(): + exhausted = package_scan().modules.exhausted + assert not exhausted, ( + f"{len(exhausted)} name lookups nested deeper than MAX_RESOLUTION_DEPTH " + "and were cut, so what they found may depend on the order modules " + "were scanned in. Raise the limit. One cut: " + repr(next(iter(exhausted))) + ) + + +def test_ratchet_only_shrinks(): + total = sum(sum(entries.values()) for entries in RATCHET.values()) + assert total <= RATCHET_CEILING, ( + "RATCHET lists engine-time draws still to move into the data. Draw " + "a new stochastic input in the dataset build instead of adding it." + ) + + +def _is_offline(name: str, offline=OFFLINE_MODULES) -> bool: + return _matches(name, offline) + + +# An offline package the import check tests against; the package has none. +EXAMPLE_OFFLINE_MODULES = {"policyengine_uk.tools.offline"} + + +@pytest.mark.parametrize( + "source", + [ + "import policyengine_uk.tools.offline.generate_tests\n", + "from policyengine_uk.tools.offline import generate_tests\n", + "from .offline import generate_tests\n", + "import importlib\n" + "importlib.import_module('policyengine_uk.tools.offline.generate_tests')\n", + "from importlib import import_module as load\n" + "load('.offline', package='policyengine_uk.tools')\n", + "__import__('policyengine_uk.tools.offline')\n", + "import importlib\n" + "importlib.import_module(name='policyengine_uk.tools.offline')\n", + "import importlib\nimportlib.import_module('.offline', 'policyengine_uk.tools')\n", + "__import__('policyengine_uk.tools', fromlist=['offline'])\n", + "__import__('offline', globals(), locals(), [], 1)\n", + ], +) +def test_offline_import_check_sees_each_import_form(source): + imports = analyse(source, package="policyengine_uk.tools").imports + assert any(_is_offline(name, EXAMPLE_OFFLINE_MODULES) for _, name in imports) + + +@pytest.mark.parametrize( + "source, offline", + [ + ( + "from importlib import import_module\nimport_module('..offline', __name__)\n", + True, + ), + ("__import__('offline', {'__name__': __name__}, None, [], 1)\n", True), + # helper is a module, not a package: Python looks for helper.offline. + ( + "from importlib import import_module\nimport_module('.offline', __name__)\n", + False, + ), + ], +) +def test_offline_import_check_resolves_against_the_modules_name(source, offline): + imports = analyse( + source, + module="policyengine_uk.tools.helper", + package="policyengine_uk.tools", + ).imports + assert ( + any(_is_offline(name, EXAMPLE_OFFLINE_MODULES) for _, name in imports) + is offline + ) + + +def test_model_never_imports_offline_modules(): + importers = [ + f"{relative}:{line}" + for relative, module, _, _ in package_modules() + if not _is_offline(module) + for line, name in package_scan().imports[relative] + if _is_offline(name) + ] + assert not importers, ( + "Offline modules are exempt from the draw check only because the " + "model never imports them:\n" + "\n".join(importers) + ) From 6f4d3313060987611bd00da46d3f392ffda1013e Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 5 Nov 2026 01:35:35 -0500 Subject: [PATCH 02/11] Correct the guard's comments and add two sys.modules test rows The guard's docstring now says it looks for draws and not for the clock (#2875). The comments on the splitmix64 and Universal Credit entries say which constants the scan recognises, which entries reach the hash and how Simulation.__init__ applies the modifier. Two test rows follow PolicyEngine/policyengine-us#10447 at 221f2fe44, and a changelog fragment records the guard. Co-Authored-By: Claude Opus 5.5 --- changelog.d/2874.changed.md | 1 + .../test_no_engine_time_randomness.py | 43 +++++++++++-------- 2 files changed, 26 insertions(+), 18 deletions(-) create mode 100644 changelog.d/2874.changed.md diff --git a/changelog.d/2874.changed.md b/changelog.d/2874.changed.md new file mode 100644 index 0000000000..aaf800f40b --- /dev/null +++ b/changelog.d/2874.changed.md @@ -0,0 +1 @@ +The test suite fails on any new random draw, sampling call or hash of entity ids in the package, and lists today's draws in a ratchet. diff --git a/policyengine_uk/tests/code_health/test_no_engine_time_randomness.py b/policyengine_uk/tests/code_health/test_no_engine_time_randomness.py index 5b38e94666..6c60b15e90 100644 --- a/policyengine_uk/tests/code_health/test_no_engine_time_randomness.py +++ b/policyengine_uk/tests/code_health/test_no_engine_time_randomness.py @@ -1,18 +1,18 @@ """The model draws nothing while a simulation is built or run. A formula is a pure, deterministic function of its inputs -(policyengine-core#500), and seeding does not exempt a draw. Stochastic -inputs such as take-up flags are drawn when the dataset is built and stored -as input variables; the model compares them with its rules. policyengine-core#519 checks formulas -for numpy and stdlib randomness when variables register, but a hand-written -integer hash of entity ids passes that check, so this scan looks for the -hash as well. +(policyengine-core#500), and seeding does not exempt a draw. The dataset +build draws stochastic inputs such as take-up flags and stores them as input +variables; the model compares them with its rules. policyengine-core#519 +checks formulas for numpy and stdlib randomness when variables register, but +a hand-written integer hash of entity ids passes that check, so this scan +looks for the hash as well. Every module of the package except tests is parsed (formulas, reforms, -scenarios, simulation code, data loading, uprating and utilities). A name resolves to -everything it may refer to where it is used: the targets of every binding -that may be the one read, in each scope Python looks in, so the scan errs -toward reporting. Bindings come from imports, named imports and re-exports +scenarios, simulation code, data loading, uprating and utilities). A name +resolves to everything it may refer to where it is used: the targets of every +binding that may be the one read, in each scope Python looks in, so the scan +errs toward reporting. Bindings come from imports, named imports and re-exports followed across the package, star imports, assignments of a name or attribute chain (``xp = np``, also in tuples), default arguments, ``global`` and ``nonlocal`` (the nearest enclosing function that binds the name, @@ -144,8 +144,9 @@ class body it is nested in binds, or its ``global`` and ``nonlocal`` the values calls return (``r = importlib.import_module("random")`` is caught at the import; a later ``r.random()`` is not), names ``match`` patterns capture, or methods as draw -helpers. A lookup that nests deeper than MAX_RESOLUTION_DEPTH fails the scan -rather than passing unresolved. +helpers. It looks for draws, not for other reads that differ between runs, +such as the clock (policyengine-uk#2875). A lookup that nests deeper than +MAX_RESOLUTION_DEPTH fails the scan rather than passing unresolved. EXEMPT lists reviewed occurrences that are not draws; RATCHET lists today's engine-time draws that still have to move into the data. Both count exact @@ -170,9 +171,10 @@ class body it is nested in binds, or its ``global`` and ``nonlocal`` TESTS = PACKAGE / "tests" RATCHET = { - # The splitmix64 finalizer that maps entity ids to uniform draws on - # [0, 1): its three mixing constants and three xorshift steps. Every - # other entry calls it, or calls what calls it. + # The splitmix64 finaliser that maps entity ids to uniform draws on + # [0, 1): the three splitmix64 constants the scan recognises and three + # xorshift steps. The formula, uprating and simulation loading entries + # below call it, or call what calls it. "utils/stochastic.py": { ("splitmix64_uniform", "hash constant"): 3, ("splitmix64_uniform", "xorshift"): 3, @@ -219,9 +221,11 @@ class body it is nested in binds, or its ``global`` and ``nonlocal`` "dynamics/participation.py": { ("apply_participation_responses", "random source"): 1, }, - # UC rebalancing scenario: a seeded numpy Generator draws one value per - # benefit unit and compares it with the new claimant share; the Scenario - # at module level passes on the modifier. + # UC rebalancing: a seeded numpy Generator draws one value per benefit + # unit and compares it with the new claimant share. Simulation.__init__ + # applies it to every simulation through the Scenario object, a call the + # scan does not follow; the Scenario at module level passes on the + # modifier. "scenarios/uc_reform.py": { ("add_universal_credit_reform", "random source"): 1, ("", "draw helper call"): 1, @@ -3777,6 +3781,9 @@ def test_a_computed_all_may_name_any_module_below_its_package(): ("import pkg\ng = getattr(pkg, 'a')\n", {"pkg.a"}, set()), ("import sys\nsys.modules['pkg.a'].f()\n", {"pkg.a"}, set()), ("from sys import modules\nmodules.get('pkg.a').f()\n", {"pkg.a"}, set()), + # sys.modules holds the module itself, not its top package. + ("import sys\nsys.modules['pkg.sub'].a.f()\n", {"pkg.sub.a"}, set()), + ("import sys\nsys.modules.get('pkg.sub').a.f()\n", {"pkg.sub.a"}, set()), ("f = lambda: __import__('pkg').a.f()\n", set(), {"pkg.a"}), # __import__ returns the top package; import_module the module. ("__import__('pkg.sub').a.f()\n", {"pkg.a", "pkg.sub.a"}, set()), From 9465138688471d554e3cc3a85776e60ef03fec67 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 5 Nov 2026 06:56:15 +0000 Subject: [PATCH 03/11] Register stored-draw inputs and count a recorded input only where core reads it Define uc_new_claimant_health_element_draw, pip_phase_in_draw and would_claim_student_loan as yearly inputs with no formula, and register them with the supplied-input helpers. supplied_input_periods, and with it supplied_by, latest_supplied_input_period and supplied_input, now count a recorded key only when its branch is the one core reads for the period and the tier core reads there still holds the input. latest_supplied_input_period ranks with core's _latest_input_key. set_model_value refuses the three new columns. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LwQq6HUE9hj2BPEjEXLL68 --- policyengine_uk/tests/stored_draw_fixtures.py | 65 ++++++ .../tests/test_supplied_input_tightening.py | 203 ++++++++++++++++++ policyengine_uk/utils/supplied_inputs.py | 160 ++++++++++---- .../gov/dwp/pip/pip_phase_in_draw.py | 17 ++ .../uc_new_claimant_health_element_draw.py | 19 ++ .../student_loans/would_claim_student_loan.py | 17 ++ 6 files changed, 443 insertions(+), 38 deletions(-) create mode 100644 policyengine_uk/tests/stored_draw_fixtures.py create mode 100644 policyengine_uk/tests/test_supplied_input_tightening.py create mode 100644 policyengine_uk/variables/gov/dwp/pip/pip_phase_in_draw.py create mode 100644 policyengine_uk/variables/gov/dwp/universal_credit/disability_element/limited_work_ability/uc_new_claimant_health_element_draw.py create mode 100644 policyengine_uk/variables/gov/hmrc/student_loans/would_claim_student_loan.py diff --git a/policyengine_uk/tests/stored_draw_fixtures.py b/policyengine_uk/tests/stored_draw_fixtures.py new file mode 100644 index 0000000000..165bb7949b --- /dev/null +++ b/policyengine_uk/tests/stored_draw_fixtures.py @@ -0,0 +1,65 @@ +"""Hand-built households and helpers for the stored-draw tests.""" + +import numpy as np +import pandas as pd + +F64 = np.float64 +F32 = np.float32 + +ID_COLUMNS = ( + "person_id", + "person_benunit_id", + "person_household_id", + "benunit_id", + "household_id", +) + + +def one_person_units_frame(count=3, years=(2025,), columns=None, ids=None): + """A frame of ``count`` adults, each alone in a benefit unit and a + household, with identifiers and ``columns`` (name -> array, or a dict of + year -> array) in each year.""" + ids = np.arange(1, count + 1) if ids is None else np.asarray(ids) + data = {} + for year in years: + for variable in ID_COLUMNS: + data[f"{variable}__{year}"] = ids + data[f"age__{year}"] = np.full(count, 30) + data[f"household_weight__{year}"] = np.ones(count) + for name, values in (columns or {}).items(): + by_year = values if isinstance(values, dict) else {years[0]: values} + for year, array in by_year.items(): + data[f"{name}__{year}"] = array + return pd.DataFrame(data) + + +def encode(x, thresholds): + """The plan's reference encoder for decision-preserving float32 draws + (tools/encode_float32_draws.py), copied for the tests: returns (y, + feasible, adjusted) for legacy float64 draws x and covered thresholds.""" + x = np.asarray(x, dtype=F64) + t = np.unique(np.asarray(thresholds, dtype=F64)) + n_le = np.searchsorted(t, x, side="right") # thresholds <= x + lo = np.where(n_le > 0, t[np.maximum(n_le - 1, 0)], 0.0) + hi = np.where(n_le < len(t), t[np.minimum(n_le, len(t) - 1)], 1.0) + lo = np.maximum(lo, 0.0) + hi = np.minimum(hi, 1.0) + + y = x.astype(F32) + yd = y.astype(F64) + ok = (yd >= lo) & (yd < hi) + + up = lo.astype(F32) + up = np.where(up.astype(F64) < lo, np.nextafter(up, F32(np.inf)), up) + down = hi.astype(F32) + down = np.where(down.astype(F64) >= hi, np.nextafter(down, F32(-np.inf)), down) + + y = np.where(ok, y, np.where(yd < lo, up, down)).astype(F32) + yd = y.astype(F64) + feasible = (yd >= lo) & (yd < hi) + return y, feasible, ~ok + + +# Thresholds the consumers compare stored draws with. +PIP_THRESHOLDS = [0.25 * (year - 2025) / 4 for year in range(2025, 2030)] +SHARE_GRID = sorted({float(f"{k / 100:.2f}") for k in range(0, 101)}) diff --git a/policyengine_uk/tests/test_supplied_input_tightening.py b/policyengine_uk/tests/test_supplied_input_tightening.py new file mode 100644 index 0000000000..04ef75178d --- /dev/null +++ b/policyengine_uk/tests/test_supplied_input_tightening.py @@ -0,0 +1,203 @@ +"""Which recorded inputs the supplied-input helpers count. + +Invariants: + +1. Registration: the three stored inputs (uc_new_claimant_health_element_draw, + pip_phase_in_draw, would_claim_student_loan) are registered, so the helpers + answer for them. +2. Ranking: latest_supplied_input_period picks, of the supplied periods that + start before the request, the one policyengine-core's carry-over order + (_latest_input_key) ranks last, whatever order they were stored in. +3. A recorded key counts only where the engine reads it: its branch is the + first readable branch that stores the period, and the tier read there + (memory before disk) still holds the input. A value cached over the input + in the same slot, a child branch's value over a parent's input and a + memory cache over an input on disk each leave the key uncounted, while the + input it shadows is still recorded. +4. set_model_value refuses the three stored inputs. +""" + +import numpy as np +import pytest +from hypothesis import HealthCheck, given, settings +from hypothesis import strategies as st +from policyengine_core.data_storage import OnDiskStorage +from policyengine_core.periods import period +from policyengine_core.simulations.simulation import _latest_input_key + +from policyengine_uk import Simulation +from policyengine_uk.tests.stored_draw_fixtures import one_person_units_frame +from policyengine_uk.utils.supplied_inputs import ( + REQUIRED_STORED_INPUTS, + SUPPLIED_INPUT_VARIABLES, + latest_supplied_input_period, + set_model_value, + supplied_by, + supplied_input, + supplied_input_periods, +) + +NEW_COLUMNS = ( + "uc_new_claimant_health_element_draw", + "pip_phase_in_draw", + "would_claim_student_loan", +) +DRAW = "pip_phase_in_draw" +# A registered input with no formula, whose readers ask the helpers. +OTHER = "adjusted_net_income" + + +def built(): + return Simulation(dataset=one_person_units_frame(count=1)) + + +@pytest.fixture(scope="module") +def shared(): + return built() + + +def periods_of(sim, name=DRAW): + return supplied_input_periods(sim.get_variable_population(name), name) + + +def test_new_columns_are_registered_and_required(): + assert set(NEW_COLUMNS) <= SUPPLIED_INPUT_VARIABLES + assert set(NEW_COLUMNS) == REQUIRED_STORED_INPUTS + sim = built() + for name in NEW_COLUMNS: + population = sim.get_variable_population(name) + assert supplied_input_periods(population, name) == [] + assert not supplied_by(population, name, period(2025)) + + +@pytest.mark.parametrize("name", NEW_COLUMNS) +def test_set_model_value_refuses_a_stored_input(name): + sim = built() + value = np.array([True]) if name == "would_claim_student_loan" else [0.5] + with pytest.raises(ValueError, match=name): + set_model_value(sim, name, 2025, value) + assert periods_of(sim, name) == [] + assert sim.get_holder(name).get_known_periods() == [] + + +def test_set_model_value_still_stores_other_variables(): + sim = built() + set_model_value(sim, OTHER, 2025, np.array([10.0])) + assert ( + supplied_input(sim.get_variable_population(OTHER), OTHER, period(2025)) is None + ) + + +def test_two_periods_starting_on_one_day_rank_by_core_order(): + for order in ( + [period("year:2024:2"), period(2024)], + [period(2024), period("year:2024:2")], + ): + sim = built() + for stored in order: + sim.set_input(DRAW, stored, [0.5]) + population = sim.get_variable_population(DRAW) + # Both start on 1 January 2024; core carries the one that ends last. + assert latest_supplied_input_period(population, DRAW, period(2026)) == period( + "year:2024:2" + ) + + +@settings( + max_examples=25, + deadline=None, + derandomize=True, + suppress_health_check=[HealthCheck.too_slow, HealthCheck.function_scoped_fixture], +) +@given( + stored=st.lists( + st.tuples(st.integers(2020, 2026), st.integers(1, 3)), + min_size=1, + max_size=4, + unique=True, + ), + before=st.integers(2020, 2028), +) +def test_latest_supplied_period_matches_core_carry_order(shared, stored, before): + sim = shared + sim.delete_arrays(DRAW) + for start, size in stored: + sim.set_input(DRAW, period(f"year:{start}:{size}"), [0.5]) + population = sim.get_variable_population(DRAW) + earlier = [ + period(f"year:{start}:{size}") for start, size in stored if start < before + ] + expected = max(earlier, key=lambda p: _latest_input_key(p, "year"), default=None) + assert latest_supplied_input_period(population, DRAW, period(before)) == expected + sim.delete_arrays(DRAW) + + +def test_a_value_cached_over_an_input_in_its_slot_is_not_supplied(): + sim = built() + sim.set_input(OTHER, 2025, [100]) + population = sim.get_variable_population(OTHER) + assert supplied_input_periods(population, OTHER) == [period(2025)] + holder = sim.get_holder(OTHER) + # A low-level cache write: core drops the slot's input tier but keeps the + # key in its record of inputs. + holder.put_in_cache(np.array([999.0]), period(2025), "default", derived=False) + assert (OTHER, "default", period(2025)) in sim._user_input_keys + assert supplied_input_periods(population, OTHER) == [] + assert supplied_input(population, OTHER, period(2025)) is None + assert not supplied_by(population, OTHER, period(2026)) + + +def test_a_child_branch_value_over_a_parent_input_is_not_supplied_on_the_child(): + sim = built() + sim.set_input(OTHER, 2025, [100]) + child = sim.get_branch("child") + child_holder = child.get_holder(OTHER) + child_holder.put_in_cache(np.array([999.0]), period(2025), "child", derived=False) + # Core reads the child's value from the child, the parent's input from the parent. + assert child_holder.get_array(period(2025), "child").tolist() == [999] + assert sim.get_holder(OTHER).get_array(period(2025), "default").tolist() == [100] + child_population = child.get_variable_population(OTHER) + assert supplied_input_periods(child_population, OTHER) == [] + assert supplied_input(child_population, OTHER, period(2025)) is None + parent_population = sim.get_variable_population(OTHER) + assert supplied_input_periods(parent_population, OTHER) == [period(2025)] + np.testing.assert_array_equal( + supplied_input(parent_population, OTHER, period(2025)), [100] + ) + + +def test_a_parent_input_with_nothing_on_the_child_is_supplied_on_the_child(): + sim = built() + sim.set_input(OTHER, 2025, [100]) + child = sim.get_branch("child") + population = child.get_variable_population(OTHER) + assert supplied_input_periods(population, OTHER) == [period(2025)] + np.testing.assert_array_equal( + supplied_input(population, OTHER, period(2025)), [100] + ) + + +@pytest.mark.parametrize("memory_cache", [False, True]) +def test_an_input_on_disk_counts_only_when_memory_does_not_hide_it( + memory_cache, tmp_path +): + sim = built() + sim.set_input(OTHER, 2025, [100]) + holder = sim.get_holder(OTHER) + source = period(2025) + holder._memory_storage.delete(source, "default") + holder._disk_storage = OnDiskStorage(str(tmp_path)) + holder._disk_storage.put(np.array([100.0]), source, "default") + holder._user_input_storage[("default", str(source))] = ( + source, + frozenset({"disk"}), + ) + population = sim.get_variable_population(OTHER) + assert supplied_input_periods(population, OTHER) == [source] + if memory_cache: + holder._memory_storage.put(np.array([999.0]), source, "default", derived=True) + # Core reads memory first, so the cache, not the input on disk. + assert supplied_input_periods(population, OTHER) == [] + assert supplied_input(population, OTHER, source) is None + else: + np.testing.assert_array_equal(supplied_input(population, OTHER, source), [100]) diff --git a/policyengine_uk/utils/supplied_inputs.py b/policyengine_uk/utils/supplied_inputs.py index d8a01bb8c8..936f5823ad 100644 --- a/policyengine_uk/utils/supplied_inputs.py +++ b/policyengine_uk/utils/supplied_inputs.py @@ -23,6 +23,15 @@ those variables. ``Simulation.delete_arrays`` forgets such records too, for any variable it deletes. +A recorded input counts only where the engine reads it. For a period, the +engine reads the first branch, of this one and those it was made from, that +stores a value, and in that branch memory before disk +(``Holder._branch_storing``, ``Holder._get_array_from_storage``). So a key +counts only when its branch is that branch and the tier the engine reads +there still holds the input (``_supplied_value_storage``). A value cached over +the input in the same tier, a child branch's value over a parent's input, and +a value cached in memory over an input on disk each leave the key uncounted. + A scenario that stores a value the model calculated (``set_model_value``) records it as the model's, and the helpers do not return it. The record of such a value lasts only while it is the value stored: ``Simulation`` forgets @@ -40,6 +49,7 @@ from policyengine_core.holders.helpers import _is_input as _stored_as_input_in from policyengine_core.holders.helpers import get_stored_input from policyengine_core.periods import Period +from policyengine_core.simulations.simulation import _latest_input_key # Variables that formulas ask about through the helpers below. @@ -90,6 +100,22 @@ "uc_carer_element", "uc_LCWRA_element", "universal_childcare_entitlement", + # Stored inputs the dataset build assigns. Their readers ask whether a + # dataset supplied them before they read them. + "pip_phase_in_draw", + "uc_new_claimant_health_element_draw", + "would_claim_student_loan", + } +) + +# Stored inputs only a dataset or a user supplies: set_model_value refuses +# them, so the model never stores its own value where the engine would read +# one of these. +REQUIRED_STORED_INPUTS = frozenset( + { + "pip_phase_in_draw", + "uc_new_claimant_health_element_draw", + "would_claim_student_loan", } ) @@ -134,9 +160,53 @@ def drop_missing_supplied_inputs(simulation, variable_name: str) -> None: ) +def _supplied_value_storage(holder, period: Period, branch_name: str): + """The storage tier (memory or disk) the engine reads for ``period`` on + ``branch_name``, if it still holds the input stored there; otherwise + None. + + The engine reads memory before disk. That tier must hold a value the + engine did not calculate, and the holder's record of where the input went + (``Holder._user_input_storage``) must list it, as + ``policyengine_core.holders.helpers._is_input`` decides. A slot restored + with no such record uses its tier's mark instead. An input on disk under a + value cached in memory does not count: the engine reads the cache. + """ + for storage_name, storage in ( + ("memory", holder._memory_storage), + ("disk", holder._disk_storage), + ): + if storage is not None and storage.has(period, branch_name): + if _stored_as_input_in(holder, period, branch_name, storage_name): + return storage + return None + return None + + +def stored_supplied_value(holder, period: Period, branch_name: str): + """The input ``holder`` stores for ``period`` on ``branch_name``, from the + tier the engine reads, or None (see ``_supplied_value_storage``).""" + storage = _supplied_value_storage(holder, period, branch_name) + return None if storage is None else storage.get(period, branch_name) + + +def _counted_input(simulation, holder, key) -> Optional[np.ndarray]: + """The value of the recorded input ``key`` if the engine reads it for its + period from this simulation's branch, and it is not the model's own + value (``set_model_value``); otherwise None.""" + name, branch_name, period = key + if holder._branch_storing(period, simulation.branch_name) != branch_name: + return None + value = stored_supplied_value(holder, period, branch_name) + if value is None or _is_model_value(simulation, key, value): + return None + return value + + def supplied_input_periods(population, variable_name: str) -> List[Period]: """Periods for which ``variable_name`` was supplied as an input, on this - branch or one it was made from, earliest first.""" + branch or one it was made from, earliest first. A period counts only + where the engine reads that input for it (see the module docstring).""" _check_registered(variable_name) simulation = population.simulation branch_names = set(_visible_branch_names(simulation)) @@ -144,16 +214,11 @@ def supplied_input_periods(population, variable_name: str) -> List[Period]: holder = population.get_holder(variable_name) return sorted( { - period - for name, branch_name, period in input_keys - if name == variable_name - and branch_name in branch_names - and holder._get_array_from_storage(period, branch_name) is not None - and not _is_model_value( - simulation, - (name, branch_name, period), - holder._get_array_from_storage(period, branch_name), - ) + key[2] + for key in list(input_keys) + if key[0] == variable_name + and key[1] in branch_names + and _counted_input(simulation, holder, key) is not None }, key=lambda period: period.start, ) @@ -162,14 +227,26 @@ def supplied_input_periods(population, variable_name: str) -> List[Period]: def latest_supplied_input_period( population, variable_name: str, before: Period ) -> Optional[Period]: - """The latest period before ``before`` for which ``variable_name`` was - supplied as an input (see ``supplied_input_periods``), or None.""" + """The latest period starting before ``before`` for which + ``variable_name`` was supplied as an input (see + ``supplied_input_periods``), or None. policyengine-core's carry-over + order picks the latest (``_latest_input_key``): a period of the + variable's own unit before one of another unit, then the one that starts + last, then the one that ends last, so two periods that start on the same + day never leave the choice to the order they were stored in.""" + definition_period = population.simulation.tax_benefit_system.get_variable( + variable_name + ).definition_period earlier = [ period for period in supplied_input_periods(population, variable_name) if period.start < before.start ] - return earlier[-1] if earlier else None + return max( + earlier, + key=lambda period: _latest_input_key(period, definition_period), + default=None, + ) def supplied_by(population, variable_name: str, period: Period) -> bool: @@ -215,29 +292,27 @@ def supplied_input( simulation = population.simulation input_keys = getattr(simulation, "_user_input_keys", None) or () holder = population.get_holder(variable_name) - for branch_name in _visible_branch_names(simulation): - if (variable_name, branch_name, period) in input_keys: - # Read the stored input itself: Holder.get_array would fall back - # to other branches' values, including cached calculations. - # A record can outlive its stored array (see above); a missing - # stored array must be ignored even if its key remains. - value = holder._get_array_from_storage(period, branch_name) - if value is not None: - # A value the model stored itself (set_model_value) is not - # supplied, though the engine keeps it like an input. - if _is_model_value( - simulation, (variable_name, branch_name, period), value - ): - return None - # A neutralized variable reads as its default, as it does - # through the engine. (PolicyEngine-UK builds no - # gov.abolitions parameters, so there is no abolition switch - # to honour here.) - variable = simulation.tax_benefit_system.get_variable(variable_name) - if variable.is_neutralized: - return holder.default_array() - return value - return None + # The branch the engine reads for the period: this one, else the nearest + # it was made from that stores a value. Holder.get_array would also fall + # back to it; read the stored input itself, from the tier the engine + # reads, so a cached value is never taken for one. + branch_name = holder._branch_storing(period, simulation.branch_name) + key = (variable_name, branch_name, period) + if branch_name is None or key not in input_keys: + return None + value = _counted_input(simulation, holder, key) + if value is None: + # A value the model stored itself (set_model_value) is not supplied, + # though the engine keeps it like an input; nor is a value cached + # over the input. + return None + # A neutralized variable reads as its default, as it does through the + # engine. (PolicyEngine-UK builds no gov.abolitions parameters, so there + # is no abolition switch to honour here.) + variable = simulation.tax_benefit_system.get_variable(variable_name) + if variable.is_neutralized: + return holder.default_array() + return value def _is_model_value(simulation, key, value) -> bool: @@ -261,8 +336,17 @@ def set_model_value(simulation, variable_name: str, period, value) -> None: do not hold it (see ``forget_replaced_model_values``, ``drop_missing_model_values`` and ``keep_population_records``). If the value was supplied before (a scenario adjusting a user's entry), the new - value stays supplied.""" + value stays supplied. + + A variable in ``REQUIRED_STORED_INPUTS`` holds only what a dataset or a + user supplies, so this refuses it.""" _check_registered(variable_name) + if variable_name in REQUIRED_STORED_INPUTS: + raise ValueError( + f"{variable_name} holds only values a dataset or a user supplies, " + "so set_model_value refuses it. Store a supplied value with " + "Simulation.set_input." + ) period = periods.period(period) population = simulation.get_variable_population(variable_name) already_supplied = supplied_input(population, variable_name, period) is not None diff --git a/policyengine_uk/variables/gov/dwp/pip/pip_phase_in_draw.py b/policyengine_uk/variables/gov/dwp/pip/pip_phase_in_draw.py new file mode 100644 index 0000000000..fbdabd17b4 --- /dev/null +++ b/policyengine_uk/variables/gov/dwp/pip/pip_phase_in_draw.py @@ -0,0 +1,17 @@ +from policyengine_uk.model_api import * + + +class pip_phase_in_draw(Variable): + value_type = float + entity = Person + label = "PIP phase-in scenario draw" + documentation = ( + "A uniform in [0, 1) that the dataset build assigns to each person " + "and stores. It is an allocation device, not an observed fact: the " + "opt-in PIP phase-in scenario removes a person's PIP in a year from " + "2025 to 2029 when this draw falls below that year's phase-in share. " + "Its situation default, 1.0, falls below no share. A simulation " + "without the column keeps the scenario's own seeded draw." + ) + definition_period = YEAR + default_value = 1.0 diff --git a/policyengine_uk/variables/gov/dwp/universal_credit/disability_element/limited_work_ability/uc_new_claimant_health_element_draw.py b/policyengine_uk/variables/gov/dwp/universal_credit/disability_element/limited_work_ability/uc_new_claimant_health_element_draw.py new file mode 100644 index 0000000000..7b4e3f254c --- /dev/null +++ b/policyengine_uk/variables/gov/dwp/universal_credit/disability_element/limited_work_ability/uc_new_claimant_health_element_draw.py @@ -0,0 +1,19 @@ +from policyengine_uk.model_api import * + + +class uc_new_claimant_health_element_draw(Variable): + value_type = float + entity = BenUnit + label = "Universal Credit new-claimant health element draw" + documentation = ( + "A uniform in [0, 1) that the dataset build assigns to each benefit " + "unit and stores. It is an allocation device, not an observed fact: " + "in each Universal Credit rebalancing year, a simulation built from " + "data that supplies no new-claimant status for the year gives a " + "benefit unit the new-claimant LCWRA rate when this draw falls below " + "gov.dwp.universal_credit.rebalancing.new_claimant_share. Its " + "situation default, 1.0, falls below no share. A dataset without " + "the column keeps the model's own seeded draw." + ) + definition_period = YEAR + default_value = 1.0 diff --git a/policyengine_uk/variables/gov/hmrc/student_loans/would_claim_student_loan.py b/policyengine_uk/variables/gov/hmrc/student_loans/would_claim_student_loan.py new file mode 100644 index 0000000000..331b82553b --- /dev/null +++ b/policyengine_uk/variables/gov/hmrc/student_loans/would_claim_student_loan.py @@ -0,0 +1,17 @@ +from policyengine_uk.model_api import * + + +class would_claim_student_loan(Variable): + value_type = bool + entity = Person + label = "Would take up a student loan if eligible" + documentation = ( + "Whether a tertiary-educated person with no recorded student loan " + "takes one up when the projection of later years makes them a new " + "borrower. The dataset build assigns it once per person and stores " + "it. It is a behavioural assumption, not an observed fact: the build " + "applies the model's graduate take-up rate. Its situation default is " + "false. A dataset without the column keeps the model's own draw." + ) + definition_period = YEAR + default_value = False From 23d9121c7ab04befb6539602c3648036802695b0 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 5 Nov 2026 07:00:28 +0000 Subject: [PATCH 04/11] Check every read of a stored draw against the input core reads input_source(population, name, period) names the supplied input policyengine-core reads for a registered draw, returns None where none reaches the period, and raises for anything else. In a simulation built from data, Simulation.calculate applies it before core, drops the request's fast-cache entry, and checks core's answer against the source after. calculate_add, calculate_divide and calculate_output refuse a draw, and each counts as a running calculation. set_input refuses a new draw while a calculation runs or a derived value of the draw is stored, and checks would_claim_student_loan's raw values before core casts them. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LwQq6HUE9hj2BPEjEXLL68 --- policyengine_uk/simulation.py | 76 ++- .../tests/test_stored_draw_input_source.py | 613 ++++++++++++++++++ policyengine_uk/utils/stored_draws.py | 414 ++++++++++++ 3 files changed, 1100 insertions(+), 3 deletions(-) create mode 100644 policyengine_uk/tests/test_stored_draw_input_source.py create mode 100644 policyengine_uk/utils/stored_draws.py diff --git a/policyengine_uk/simulation.py b/policyengine_uk/simulation.py index e000f0f9a3..24c762a01a 100644 --- a/policyengine_uk/simulation.py +++ b/policyengine_uk/simulation.py @@ -99,6 +99,16 @@ keep_population_records, model_value_input_records, ) +from policyengine_uk.utils.stored_draws import ( + REGISTERED_DRAW_NAMES, + STUDENT_LOAN_TAKE_UP, + calculation, + check_boolean_input, + check_draw_answer, + check_new_draw_input, + check_option_request, + input_source, +) from policyengine_uk.reforms import create_structural_reforms_from_parameters from policyengine_uk.parameters.gov.simulation.labour_supply_responses.aliases import ( canonicalize_parameter_changes, @@ -184,6 +194,11 @@ def __reduce__(self): _input_key_logs: Dict[int, list] = {} +# Simulation.calculate: a variable that is not a stored draw read in a +# simulation built from data. +_UNCHECKED = object() + + # Cache for fully-loaded, uprated, enum-pre-encoded multi-year datasets keyed # by URL and, when selected, the CGT projection fingerprint and runtime. # Avoids repeating HDF5 reading (0.84s), uprating (0.69s) and enum @@ -634,6 +649,13 @@ def build_from_populations(self, populations) -> None: keep_population_records(self, adopted) def set_input(self, variable_name: str, period: Period, value) -> None: + if variable_name in REGISTERED_DRAW_NAMES: + # Before core stores anything: never during a calculation, and + # never over a value calculated from an earlier state. + check_new_draw_input(self, variable_name) + elif variable_name == STUDENT_LOAN_TAKE_UP: + # Before core casts the values: a missing value would load as True. + check_boolean_input(variable_name, value) # What is set here is supplied, even a value equal to the model's own, # so forget the model value (set_model_value) of each period the call # stores a value for, and only those. set_model_value records its @@ -1799,6 +1821,18 @@ def calculate( # as supplied. drop_missing_supplied_inputs(self, variable_name) + draw_source = _UNCHECKED + if variable_name in REGISTERED_DRAW_NAMES and built_from_data(self): + # A stored draw: decide which supplied input core reads, after the + # cleanups above and before core, and never answer from the fast + # cache, which can outlive the value the check examined. + draw_source = input_source( + self.get_variable_population(variable_name), variable_name, period + ) + fast_cache = getattr(self, "_fast_cache", None) + if fast_cache: + fast_cache.pop((variable_name, period), None) + outer_calculation = self._cgt_calculation_depth == 0 if outer_calculation: self._cgt_operation_inputs = set() @@ -1809,9 +1843,15 @@ def calculate( record = self._cgt_calculation_record(variable_name, period, neutralized) self._cgt_calculation_depth += 1 try: - result = super().calculate( - variable_name, period, map_to=map_to, decode_enums=decode_enums - ) + with calculation(): + if draw_source is _UNCHECKED or draw_source is None: + result = super().calculate( + variable_name, period, map_to=map_to, decode_enums=decode_enums + ) + else: + result = self._calculate_source_backed_draw( + variable_name, period, draw_source, map_to + ) if record is not None: if neutralized: self._cgt_defaulted_periods.add((variable_name, str(period))) @@ -1827,6 +1867,36 @@ def calculate( self._warn_unused_cgt_projection(period_(input_period)) return result + def _calculate_source_backed_draw(self, variable_name, period, source, map_to): + """Core's answer for a stored draw with a supplied source, checked + against that source before it is mapped or returned.""" + raw = super().calculate(variable_name, period, map_to=None, decode_enums=False) + check_draw_answer(self, variable_name, period, source, raw) + if map_to is not None: + entity = self.tax_benefit_system.get_variable(variable_name).entity.key + return self.map_result(raw, entity, map_to) + return raw + + def calculate_add(self, variable_name: str, period: Period = None, *args, **kwargs): + if variable_name in REGISTERED_DRAW_NAMES and built_from_data(self): + check_option_request(self, variable_name, "calculate_add") + with calculation(): + return super().calculate_add(variable_name, period, *args, **kwargs) + + def calculate_divide( + self, variable_name: str, period: Period = None, *args, **kwargs + ): + if variable_name in REGISTERED_DRAW_NAMES and built_from_data(self): + check_option_request(self, variable_name, "calculate_divide") + with calculation(): + return super().calculate_divide(variable_name, period, *args, **kwargs) + + def calculate_output(self, variable_name: str, period: Period = None): + if variable_name in REGISTERED_DRAW_NAMES and built_from_data(self): + check_option_request(self, variable_name, "calculate_output") + with calculation(): + return super().calculate_output(variable_name, period) + def apply_dynamics(self, year: int): """Apply dynamics to the simulation for a specific year. diff --git a/policyengine_uk/tests/test_stored_draw_input_source.py b/policyengine_uk/tests/test_stored_draw_input_source.py new file mode 100644 index 0000000000..e9c2b98794 --- /dev/null +++ b/policyengine_uk/tests/test_stored_draw_input_source.py @@ -0,0 +1,613 @@ +"""What a read of a stored draw may return, in simulations built from data. + +``input_source(population, name, period)`` names the supplied input +policyengine-core reads for a registered draw, returns None where no supplied +input reaches the period, and raises otherwise. ``Simulation.calculate`` +applies it before core and checks core's answer after. + +Invariants: + +1. Provenance: a non-derived value in the slot core reads counts only when it + is a supplied input; a derived value there only when core's carry source, + chosen by core's own order, is supplied and the cache equals it; with no + supplied source only a cached default, as absence. Anything else raises. +2. Requests: only one whole year, in storage form. Months, longer periods and + aliases raise, with an exact input, a carried one or none. +3. Definitions: the current definition and the holder's must keep the + registered entity, type and period, with no formula or other computation; + a neutralised or abolished draw, carry turned off where a read needs it, + core's macro cache and a disk tier each raise. +4. Dispatch: calculate never answers a draw from the fast cache; core's + answer must equal the source it read; ADD, DIVIDE and output refuse draws. +5. Writes: a new draw input is refused while a calculation runs or while a + derived value of the draw is stored, and accepted otherwise. +6. Differential: every accepted read returns exactly what core's own + calculate returns with the check removed and the fast-cache entry + cleared; with map_to=None it is the source's array, or with no source the + holder's default. +7. Changing nothing: input_source leaves storage, records and the fast cache + as they were. +""" + +import copy + +import numpy as np +import pytest +from policyengine_core.data_storage import OnDiskStorage +from policyengine_core.model_api import ADD, DIVIDE +from policyengine_core.periods import period +from policyengine_core.simulations import Simulation as CoreSimulation + +from policyengine_uk import Simulation +from policyengine_uk.model_api import BenUnit, Person, Variable, YEAR +from policyengine_uk.tests.stored_draw_fixtures import one_person_units_frame +from policyengine_uk.utils.stored_draws import ( + CORE_DISK_KEY_ISSUE, + InputSource, + StoredDrawError, + calculation, + input_source, +) + +DRAW = "pip_phase_in_draw" +GROUP_DRAW = "uc_new_claimant_health_element_draw" +VALUES = np.array([0.25, 0.5, 0.75], dtype=np.float32) +OTHER = np.array([0.125, 0.375, 0.625], dtype=np.float32) + + +def built(draws=None, name=DRAW): + """Three one-person units, with ``draws`` (year -> values) supplied.""" + columns = {name: draws} if draws else {} + years = tuple(sorted({2025, *(draws or {})})) + return Simulation( + dataset=one_person_units_frame(count=3, years=years, columns=columns) + ) + + +def source_of(sim, p, name=DRAW): + return input_source(sim.get_variable_population(name), name, p) + + +def core_answer(sim, p, name=DRAW, **kwargs): + """Core's own calculate, with the check removed and the request's + fast-cache entry cleared.""" + sim._fast_cache.pop((name, period(p)), None) + return CoreSimulation.calculate(sim, name, period(p), **kwargs) + + +def storage_state(sim, name=DRAW): + holder = sim.get_holder(name) + memory = holder._memory_storage + return ( + {key: np.array(value) for key, value in memory._arrays.items()}, + set(memory._inputs), + set(sim._user_input_keys), + copy.deepcopy(getattr(holder, "_user_input_storage", {})), + {key: np.array(value) for key, value in sim._fast_cache.items()}, + ) + + +def assert_same_state(before, after): + arrays, inputs, keys, locations, fast = before + arrays2, inputs2, keys2, locations2, fast2 = after + assert arrays.keys() == arrays2.keys() + for key in arrays: + np.testing.assert_array_equal(arrays[key], arrays2[key]) + assert (inputs, keys, locations) == (inputs2, keys2, locations2) + assert fast.keys() == fast2.keys() + + +# --- Provenance ------------------------------------------------------------- + + +def test_an_exact_supplied_draw_is_its_own_source_before_and_after_reads(): + sim = built({2025: VALUES}) + for _ in range(3): + source = source_of(sim, 2025) + assert isinstance(source, InputSource) + assert (source.period, source.branch_name, source.tier) == ( + period(2025), + "default", + "memory", + ) + np.testing.assert_array_equal(sim.calculate(DRAW, 2025), VALUES) + # An exact read stores nothing derived. + assert sim.get_holder(DRAW).get_known_periods() == [period(2025)] + + +def test_a_draw_supplied_for_2025_carries_to_2026_on_every_read(): + sim = built({2025: VALUES}) + first = source_of(sim, 2026) + assert first.period == period(2025) + for _ in range(3): + np.testing.assert_array_equal(sim.calculate(DRAW, 2026), VALUES) + # After the first read core caches the carried value as derived; the + # check accepts it because it equals its supplied source. + assert source_of(sim, 2026) == first + assert sim.get_holder(DRAW).is_derived(period(2026), "default") + + +def test_no_supplied_draw_is_absence_and_reads_the_default_every_time(): + sim = built() + assert source_of(sim, 2026) is None + for _ in range(2): + np.testing.assert_array_equal(sim.calculate(DRAW, 2026), [1.0, 1.0, 1.0]) + assert source_of(sim, 2026) is None + + +def test_an_input_only_for_a_later_year_leaves_earlier_years_absent(): + sim = built({2027: VALUES}) + for _ in range(2): + assert source_of(sim, 2026) is None + np.testing.assert_array_equal(sim.calculate(DRAW, 2026), [1.0, 1.0, 1.0]) + # Core returns the default without caching it. + assert period(2026) not in sim.get_holder(DRAW).get_known_periods() + + +def test_a_model_marked_value_over_an_older_supplied_input_raises(): + sim = built({2025: VALUES}) + sim.set_input(DRAW, 2026, OTHER) + holder = sim.get_holder(DRAW) + sim._model_values[(DRAW, "default", period(2026))] = np.array( + holder.get_array(period(2026), "default") + ) + with pytest.raises(StoredDrawError, match="not a supplied one"): + sim.calculate(DRAW, 2026) + + +def test_a_value_cached_as_non_derived_over_an_input_raises(): + sim = built({2025: VALUES}) + sim.get_holder(DRAW).put_in_cache(OTHER, period(2025), "default", derived=False) + with pytest.raises(StoredDrawError, match="not a supplied one"): + source_of(sim, 2025) + + +def test_a_non_input_carry_source_raises(): + sim = built() + # Non-derived, so core would carry it, but never set as an input. + sim.get_holder(DRAW).put_in_cache(OTHER, period(2024), "default", derived=False) + np.testing.assert_array_equal(core_answer(sim, 2026), OTHER) + sim._fast_cache.clear() + sim.get_holder(DRAW).delete_arrays(period(2026)) + with pytest.raises(StoredDrawError, match="not a supplied input"): + sim.calculate(DRAW, 2026) + + +def test_a_child_branch_value_over_a_parent_input_raises_on_the_child(): + sim = built({2025: VALUES}) + child = sim.get_branch("child") + child.get_holder(DRAW).put_in_cache(OTHER, period(2025), "child", derived=False) + with pytest.raises(StoredDrawError, match="branch 'child'"): + source_of(child, 2025) + assert source_of(sim, 2025).branch_name == "default" + + +def test_a_child_branch_reads_its_parent_input_and_carried_cache(): + sim = built({2025: VALUES}) + np.testing.assert_array_equal(sim.calculate(DRAW, 2026), VALUES) + child = sim.get_branch("child") + for p in (2025, 2026): + source = source_of(child, p) + assert (source.period, source.branch_name) == (period(2025), "default") + np.testing.assert_array_equal(child.calculate(DRAW, p), VALUES) + + +def test_a_clone_keeps_its_source(): + sim = built({2025: VALUES}) + sim.calculate(DRAW, 2026) + clone = sim.clone() + assert source_of(clone, 2026).period == period(2025) + np.testing.assert_array_equal(clone.calculate(DRAW, 2026), VALUES) + + +def test_a_default_cached_before_an_earlier_input_is_supplied(): + sim = built() + sim.calculate(DRAW, 2026) # caches the default as derived + # The public setter refuses while the derived default stands. + with pytest.raises(StoredDrawError, match="after the simulation has calculated"): + sim.set_input(DRAW, 2025, VALUES) + # A holder-level write gets past it; the surviving default then differs + # from the new source, so the read raises. + sim.get_holder(DRAW).set_input(period(2025), VALUES) + with pytest.raises(StoredDrawError, match="differs from its supplied source"): + sim.calculate(DRAW, 2026) + + +def test_a_stale_non_default_cache_with_no_source_raises(): + sim = built() + sim.get_holder(DRAW)._memory_storage.put( + OTHER, period(2026), "default", derived=True + ) + with pytest.raises(StoredDrawError, match="not the default"): + sim.calculate(DRAW, 2026) + + +def test_a_cached_default_is_compared_with_the_holders_own_default(): + sim = built() + sim.calculate(DRAW, 2026) + + class pip_phase_in_draw(Variable): + value_type = float + entity = Person + definition_period = YEAR + default_value = 0.5 + label = "draw with another default" + + sim.tax_benefit_system.update_variable(pip_phase_in_draw) + # The holder keeps the definition it was made with, so its default 1.0. + assert source_of(sim, 2026) is None + + +def test_input_source_changes_nothing(): + sim = built({2025: VALUES}) + sim.calculate(DRAW, 2026) + sim.calculate(DRAW, 2025) + before = storage_state(sim) + for p in (2024, 2025, 2026, 2027): + source_of(sim, p) + assert_same_state(before, storage_state(sim)) + + +# --- Root, branch and deletion --------------------------------------------- + + +def test_root_holder_replacement_raises_on_an_unequal_cache_and_keeps_an_equal_one(): + sim = built({2025: VALUES}) + sim.calculate(DRAW, 2026) + holder = sim.get_holder(DRAW) + holder.set_input(period(2025), VALUES) + np.testing.assert_array_equal(sim.calculate(DRAW, 2026), VALUES) + holder.set_input(period(2025), OTHER) + with pytest.raises(StoredDrawError, match="differs from its supplied source"): + sim.calculate(DRAW, 2026) + + +def test_branch_holder_replacement_recomputes_from_the_new_source(): + sim = built({2025: VALUES}) + branch = sim.get_branch("reform") + np.testing.assert_array_equal(branch.calculate(DRAW, 2026), VALUES) + branch.get_holder(DRAW).set_input(period(2025), OTHER, "reform") + np.testing.assert_array_equal(branch.calculate(DRAW, 2026), OTHER) + np.testing.assert_array_equal(sim.calculate(DRAW, 2026), VALUES) + + +def test_deleting_the_source_exposes_an_equal_earlier_input_or_raises(): + sim = built({2024: VALUES, 2025: VALUES}) + sim.calculate(DRAW, 2026) + sim.get_holder(DRAW).delete_arrays(period(2025)) + source = source_of(sim, 2026) + assert source.period == period(2024) + np.testing.assert_array_equal(sim.calculate(DRAW, 2026), VALUES) + + sim = built({2025: VALUES}) + sim.calculate(DRAW, 2026) + sim.get_holder(DRAW)._memory_storage.delete(period(2025), "default") + with pytest.raises(StoredDrawError, match="no supplied input"): + source_of(sim, 2026) + + +def test_a_rebuild_drops_the_fast_cache_and_a_new_holder_recomputes(): + sim = built({2025: VALUES}) + sim.calculate(DRAW, 2026) + sim._fast_cache[(DRAW, period(2026))] = OTHER + other = built({2024: OTHER}) + sim.build_from_populations(other.populations) + assert not sim._fast_cache + assert source_of(sim, 2026).period == period(2024) + np.testing.assert_array_equal(sim.calculate(DRAW, 2026), OTHER) + + +def test_holder_cache_opt_out_reads_storage_each_time(): + sim = built({2025: VALUES}) + sim.get_holder(DRAW)._do_not_store = True + for _ in range(2): + np.testing.assert_array_equal(sim.calculate(DRAW, 2026), VALUES) + sim.set_input(DRAW, 2025, OTHER) + np.testing.assert_array_equal(sim.calculate(DRAW, 2026), OTHER) + + +# --- Requests and definitions ----------------------------------------------- + + +@pytest.mark.parametrize( + "request_period", + ["2026-06", "year:2026:2", "month:2026-06:12", "year:2025-01-15"], +) +@pytest.mark.parametrize("supplied", [{2026: VALUES}, {2025: VALUES}, None]) +def test_requests_other_than_one_stored_year_raise(request_period, supplied): + sim = built(supplied) + with pytest.raises(StoredDrawError, match="read only for one year"): + sim.calculate(DRAW, request_period) + + +def test_add_divide_and_output_refuse_a_draw(): + sim = built({2025: VALUES}) + with pytest.raises(StoredDrawError, match="calculate_add"): + sim.calculate_add(DRAW, "year:2025:2") + with pytest.raises(StoredDrawError, match="calculate_divide"): + sim.calculate_divide(DRAW, "2025-06") + with pytest.raises(StoredDrawError, match="calculate_output"): + sim.calculate_output(DRAW, 2025) + population = sim.get_variable_population(DRAW) + with pytest.raises(StoredDrawError): + population(DRAW, period("2025-06"), options=[DIVIDE]) + with pytest.raises(StoredDrawError): + population(DRAW, period("year:2025:2"), options=[ADD]) + + +def test_a_neutralised_draw_raises(): + sim = built({2025: VALUES}) + sim.tax_benefit_system.neutralize_variable(DRAW) + for p in (2025, 2026): + with pytest.raises(StoredDrawError, match="neutralised"): + sim.calculate(DRAW, p) + + +def test_a_draw_abolished_by_parameter_raises(monkeypatch): + sim = built({2025: VALUES}) + system = sim.tax_benefit_system + real = system.parameters + + class Abolitions: + def __getitem__(self, name): + return name == DRAW + + class Node: + def __init__(self, inner): + self.inner = inner + + def __getattr__(self, name): + if name == "abolitions": + return Abolitions() + return getattr(self.inner, name) + + class At: + def __init__(self, inner): + self.gov = Node(inner.gov) + + def parameters(p): + return At(real(p)) + + monkeypatch.setattr(system, "parameters", parameters, raising=False) + with pytest.raises(StoredDrawError, match="abolishes"): + sim.calculate(DRAW, 2026) + + +def test_redefinitions_raise(): + for attributes in ( + {"formula": lambda person, period, parameters: 0.5}, + {"adds": ["age"]}, + {"value_type": int, "default_value": 1}, + ): + sim = built({2025: VALUES}) + body = dict( + value_type=float, + entity=Person, + definition_period=YEAR, + default_value=1.0, + label="redefined draw", + ) + body.update(attributes) + sim.tax_benefit_system.update_variable(type(DRAW, (Variable,), body)) + for p in (2025, 2026): + with pytest.raises(StoredDrawError, match="current definition"): + sim.calculate(DRAW, p) + + +def test_a_draw_moved_to_another_entity_raises(): + sim = built({2025: VALUES}) + sim.tax_benefit_system.update_variable( + type( + DRAW, + (Variable,), + dict( + value_type=float, + entity=BenUnit, + definition_period=YEAR, + default_value=1.0, + label="draw on benefit units", + ), + ) + ) + with pytest.raises(StoredDrawError): + source_of(sim, 2025) + + +def test_a_holder_kept_from_another_definition_raises(): + sim = built({2025: VALUES}) + holder = sim.get_holder(DRAW) + retained = copy.copy(holder.variable) + retained.value_type = int + holder.variable = retained + with pytest.raises(StoredDrawError, match="holder's definition"): + sim.calculate(DRAW, 2025) + + +def test_carry_off_raises_only_where_a_read_needs_it(): + sim = built({2025: VALUES}) + sim.tax_benefit_system.auto_carry_over_input_variables = False + try: + np.testing.assert_array_equal(sim.calculate(DRAW, 2025), VALUES) + np.testing.assert_array_equal(sim.calculate(DRAW, 2025), VALUES) + with pytest.raises(StoredDrawError, match="does not carry inputs over"): + sim.calculate(DRAW, 2026) + assert source_of(sim, 2024) is None + finally: + sim.tax_benefit_system.auto_carry_over_input_variables = True + + +def test_the_macro_cache_raises(): + sim = built({2025: VALUES}) + sim.macro_cache_read = True + with pytest.raises(StoredDrawError, match="macro cache"): + sim.calculate(DRAW, 2025) + + +@pytest.mark.parametrize( + "route", ["attached", "cloned_holder", "branch_with_underscore"] +) +def test_a_disk_tier_raises_naming_core_526(route, tmp_path): + sim = built({2025: VALUES}) + target = ( + sim.get_branch("with_underscore") if route == "branch_with_underscore" else sim + ) + holder = target.get_holder(DRAW) + holder._disk_storage = OnDiskStorage(str(tmp_path)) + if route == "cloned_holder": + population = target.get_variable_population(DRAW) + population._holders[DRAW] = holder.clone(population) + for p in (2025, 2026): + with pytest.raises(StoredDrawError, match=CORE_DISK_KEY_ISSUE): + target.calculate(DRAW, p) + + +# --- Dispatch --------------------------------------------------------------- + + +def test_the_fast_cache_never_answers_a_draw(): + sim = built({2025: VALUES}) + sim.calculate(DRAW, 2026) + sim._fast_cache[(DRAW, period(2026))] = OTHER + np.testing.assert_array_equal(sim.calculate(DRAW, 2026), VALUES) + + +def test_a_wrong_answer_from_core_raises(monkeypatch): + sim = built({2025: VALUES}) + real = CoreSimulation.calculate + + def wrong(self, name, *args, **kwargs): + result = real(self, name, *args, **kwargs) + return result * 0 if name == DRAW else result + + monkeypatch.setattr(CoreSimulation, "calculate", wrong) + with pytest.raises(StoredDrawError, match="not its supplied"): + sim.calculate(DRAW, 2026) + + +def test_a_spiral_default_from_core_raises(): + sim = built({2025: VALUES}) + for year in range(2010, 2010 + sim.max_spiral_loops): + sim.tracer.stack.append( + {"name": DRAW, "period": period(year), "branch_name": "default"} + ) + try: + with pytest.raises(StoredDrawError, match="not its supplied"): + sim.calculate(DRAW, 2026) + finally: + del sim.tracer.stack[:] + + +def test_tracing_and_map_to_return_what_core_returns(): + sim = built({2025: VALUES}, name=GROUP_DRAW) + sim.trace = True + for p in (2025, 2026): + for map_to in (None, "person", "household"): + ours = np.asarray(sim.calculate(GROUP_DRAW, p, map_to=map_to)) + theirs = np.asarray(core_answer(sim, p, name=GROUP_DRAW, map_to=map_to)) + np.testing.assert_array_equal(ours, theirs) + + +def test_differential_against_core_calculate(): + cases = [ + ({2025: VALUES}, [2024, 2025, 2026, 2030]), + ({2027: VALUES}, [2025, 2026, 2027, 2028]), + ({2025: VALUES, 2027: OTHER}, [2025, 2026, 2027, 2028]), + (None, [2025, 2026]), + ] + for supplied, requests in cases: + sim = built(supplied) + for p in requests * 2: + source = source_of(sim, p) + ours = np.asarray(sim.calculate(DRAW, p)) + holder = sim.get_holder(DRAW) + expected = ( + holder.default_array() + if source is None + else np.asarray( + holder._memory_storage.get(source.period, source.branch_name) + ) + ) + np.testing.assert_array_equal(ours, expected) + assert ours.dtype == np.float32 + np.testing.assert_array_equal(ours, np.asarray(core_answer(sim, p))) + + +# --- Writes ----------------------------------------------------------------- + + +def test_a_draw_cannot_be_set_while_any_calculation_runs(): + sim = built() + other = built() + with calculation(): + for target in (sim, other): + with pytest.raises(StoredDrawError, match="while a calculation"): + target.set_input(DRAW, 2025, VALUES) + assert sim.get_holder(DRAW).get_known_periods() == [] + + +def test_a_draw_set_from_inside_another_variables_formula_is_refused(): + sim = built() + + class writes_a_draw(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "writes a draw" + + def formula(person, period, parameters): + person.simulation.set_input(DRAW, period, VALUES) + return 0 + + sim.tax_benefit_system.add_variable(writes_a_draw) + with pytest.raises(StoredDrawError, match="while a calculation"): + sim.calculate("writes_a_draw", 2025) + + +@pytest.mark.parametrize("before", ["no_read", "exact_reads", "opt_out"]) +def test_a_new_draw_is_accepted_before_any_derived_value(before): + sim = built({2025: VALUES}) + if before == "exact_reads": + sim.calculate(DRAW, 2025) + sim.calculate(DRAW, 2025) + if before == "opt_out": + sim.get_holder(DRAW)._do_not_store = True + sim.calculate(DRAW, 2026) + sim.set_input(DRAW, 2025, OTHER) + np.testing.assert_array_equal(sim.calculate(DRAW, 2025), OTHER) + np.testing.assert_array_equal(sim.calculate(DRAW, 2026), OTHER) + + +def test_a_new_draw_is_refused_after_a_carried_read_until_a_reset(): + sim = built({2025: VALUES}) + sim.calculate(DRAW, 2026) + with pytest.raises(StoredDrawError, match="after the simulation has calculated"): + sim.set_input(DRAW, 2025, OTHER) + np.testing.assert_array_equal(sim.calculate(DRAW, 2026), VALUES) + sim.reset_calculations() + sim.set_input(DRAW, 2025, OTHER) + np.testing.assert_array_equal(sim.calculate(DRAW, 2026), OTHER) + + +def test_reforms_and_resets_keep_the_supplied_draw(): + sim = built({2025: VALUES}) + sim.calculate(DRAW, 2026) + sim.apply_reform({"gov.dwp.pip.daily_living.enhanced": {"2025-01-01": 200}}) + assert source_of(sim, 2026).period == period(2025) + np.testing.assert_array_equal(sim.calculate(DRAW, 2026), VALUES) + sim.reset_calculations() + np.testing.assert_array_equal(sim.calculate(DRAW, 2026), VALUES) + branch = sim.get_branch("measurement") + np.testing.assert_array_equal(branch.calculate(DRAW, 2026), VALUES) + + +def test_situations_keep_cores_path(): + sim = Simulation( + situation={ + "people": {"adult": {"age": {2025: 30}}}, + "benunits": {"benunit": {"members": ["adult"]}}, + "households": {"household": {"members": ["adult"]}}, + } + ) + np.testing.assert_array_equal( + sim.calculate_divide(DRAW, "2025-06"), np.float32([1 / 12]) + ) diff --git a/policyengine_uk/utils/stored_draws.py b/policyengine_uk/utils/stored_draws.py new file mode 100644 index 0000000000..e4dc5fadbf --- /dev/null +++ b/policyengine_uk/utils/stored_draws.py @@ -0,0 +1,414 @@ +"""Draws a dataset build stores, and the one check that decides what a read of +one may return. + +A stored draw is a uniform in [0, 1) that the producer of a dataset assigns +to each unit and stores; the model compares it with a rate it owns. The +registry below records each draw's entity, value type and definition period. +A simulation built from data reads a registered draw only through +``Simulation.calculate``, which asks ``input_source`` which input +policyengine-core will read before core computes the answer, and checks +core's answer against that input afterwards (``check_draw_answer``). + +``input_source`` changes nothing in storage and returns no array. It returns +the period, branch, storage tier and sequence number of the supplied input +core will read for the request, None where no supplied input reaches the +period, and raises otherwise. Metadata alone decides whether a value is +supplied (see ``policyengine_uk.utils.supplied_inputs``). Arrays are compared +only to accept a derived value core cached in the slot it reads: it must +equal its supplied source, or, with no supplied source, the holder's default. + +The check reads policyengine-core's private members: ``Holder._branch_storing``, +``Holder._memory_storage``, ``Holder._disk_storage``, +``Holder._stored_sequence_number`` and core's carry-over order +``_latest_input_key``. Tests pin the behaviour it relies on for each +qualified core version (see ``QUALIFIED_CORE_VERSIONS``). +""" + +from contextlib import contextmanager +from contextvars import ContextVar +from typing import NamedTuple, Optional + +import numpy as np +from policyengine_core import periods +from policyengine_core.periods import Period +from policyengine_core.simulations.simulation import _latest_input_key + +from policyengine_uk.utils.supplied_inputs import ( + SUPPLIED_INPUT_VARIABLES, + _counted_input, +) + +# The policyengine-core versions whose provenance, dispatch, lifecycle and +# loading behaviour the tests qualify. tests/code_health/test_core_version_gate.py +# fails when uv.lock or the installed core holds any other. +QUALIFIED_CORE_VERSIONS = frozenset({"3.33.3", "3.34.1"}) + +# policyengine-core splits on-disk storage keys at every underscore, so a +# branch name with one breaks its enumeration of inputs on disk. +CORE_DISK_KEY_ISSUE = "PolicyEngine/policyengine-core#526" + + +class DrawSchema(NamedTuple): + """What a registered draw's definition must keep.""" + + entity: str + value_type: type + definition_period: str + + +REGISTERED_DRAWS = { + "uc_new_claimant_health_element_draw": DrawSchema("benunit", float, periods.YEAR), + "pip_phase_in_draw": DrawSchema("person", float, periods.YEAR), +} +# Simulation.calculate tests each variable name against this set, once. +REGISTERED_DRAW_NAMES = frozenset(REGISTERED_DRAWS) +assert REGISTERED_DRAW_NAMES <= SUPPLIED_INPUT_VARIABLES + + +class InputSource(NamedTuple): + """The supplied input policyengine-core reads for a request.""" + + period: Period + branch_name: str + tier: str + sequence_number: Optional[int] + + +class StoredDrawError(ValueError): + """A read or write of a registered draw that the model refuses.""" + + +# How many pe-uk calculations (calculate, calculate_add, calculate_divide, +# calculate_output) are running in this context, on any simulation. +_calculation_depth: ContextVar[int] = ContextVar( + "policyengine_uk_calculation_depth", default=0 +) + + +@contextmanager +def calculation(): + """Count a pe-uk calculation as running while the block runs.""" + token = _calculation_depth.set(_calculation_depth.get() + 1) + try: + yield + finally: + _calculation_depth.reset(token) + + +def calculation_in_flight() -> bool: + return _calculation_depth.get() > 0 + + +def _describe(name: str, period, branch_name: Optional[str], tier=None) -> str: + where = f"{name} for {period}" + if branch_name is not None: + where += f" on branch {branch_name!r}" + if tier is not None: + where += f" ({tier})" + return where + + +def _check_schema(variable, name: str, schema: DrawSchema, which: str) -> None: + if variable is None: + raise StoredDrawError(f"{name} has no {which}.") + kept = ( + variable.entity.key == schema.entity + and variable.value_type is schema.value_type + and variable.definition_period == schema.definition_period + ) + if not kept: + raise StoredDrawError( + f"The {which} of {name} has entity {variable.entity.key!r}, type " + f"{variable.value_type.__name__} and period " + f"{variable.definition_period!r}; a stored draw needs entity " + f"{schema.entity!r}, type {schema.value_type.__name__} and period " + f"{schema.definition_period!r}." + ) + computed = [ + attribute + for attribute, value in ( + ("a formula", len(variable.formulas) > 0), + ("uprating", variable.uprating is not None), + ("defined_for", variable.defined_for is not None), + ("calculate_output", variable.calculate_output is not None), + ("adds", bool(variable.adds)), + ("subtracts", bool(variable.subtracts)), + ( + "requires_computation_after", + variable.requires_computation_after is not None, + ), + ) + if value + ] + if computed: + raise StoredDrawError( + f"The {which} of {name} has {', '.join(computed)}; a stored draw " + "is an input the dataset supplies, read as stored." + ) + + +def _check_request(holder, name: str, schema: DrawSchema, period: Period) -> None: + if ( + period.unit != schema.definition_period + or period.size != 1 + or holder._storage_period(period) != period + ): + raise StoredDrawError( + f"{name} is read only for one {schema.definition_period}, as " + f"stored; {period} is not one. A draw is a share of nothing, so " + "it has no sum over months or twelfth." + ) + + +def _check_definition(simulation, holder, name, schema, period) -> None: + system = simulation.tax_benefit_system + current = system.variables.get(name) + _check_schema(current, name, schema, "current definition") + _check_schema(holder.variable, name, schema, "holder's definition") + if current.is_neutralized or holder.variable.is_neutralized: + raise StoredDrawError( + f"{name} is neutralised, so policyengine-core would return its " + "default instead of the stored draw." + ) + try: + abolished = bool(system.parameters(period).gov.abolitions[name]) + except Exception: + # Core reads the same switch and ignores its absence. + abolished = False + if abolished: + raise StoredDrawError( + f"gov.abolitions abolishes {name} in {period}, so " + "policyengine-core would return its default instead of the " + "stored draw." + ) + if getattr(simulation, "macro_cache_read", False): + raise StoredDrawError( + f"The simulation reads policyengine-core's macro cache, which " + f"could answer {name} from a file instead of the stored draw." + ) + + +def _sequence_number(holder, period: Period, branch_name: str): + return holder._stored_sequence_number(period, branch_name) + + +def _supplied_source(simulation, holder, name, period, branch_name): + """``InputSource`` for the input stored for ``period`` on ``branch_name`` + if it counts as supplied there, else None.""" + if _counted_input(simulation, holder, (name, branch_name, period)) is None: + return None + return InputSource( + period, branch_name, "memory", _sequence_number(holder, period, branch_name) + ) + + +def _carry_source(simulation, holder, name, schema, period) -> Optional[InputSource]: + """The input core's carry-over would read for ``period``: of the periods + ``Holder.get_input_periods`` lists that start no later than it, the one + ``_latest_input_key`` ranks last. It must be supplied.""" + branch_name = simulation.branch_name + candidates = [ + input_period + for input_period in holder.get_input_periods(branch_name) + if input_period.start <= period.start + ] + if not candidates: + return None + source_period = max( + candidates, + key=lambda candidate: _latest_input_key(candidate, schema.definition_period), + ) + source_branch = holder._branch_storing(source_period, branch_name) + source = _supplied_source(simulation, holder, name, source_period, source_branch) + if source is None: + raise StoredDrawError( + f"policyengine-core would carry " + f"{_describe(name, source_period, source_branch, 'memory')} into " + f"{period}, and that value is not a supplied input." + ) + return source + + +def _source_array(holder, source: InputSource) -> np.ndarray: + stored = holder._memory_storage.get(source.period, source.branch_name) + return np.asarray(stored).astype(holder.variable.dtype) + + +def _same_array(first, second) -> bool: + first = np.asarray(first) + second = np.asarray(second) + return ( + first.dtype == second.dtype + and first.shape == second.shape + and bool(np.array_equal(first, second)) + ) + + +def input_source(population, name: str, period) -> Optional[InputSource]: + """The supplied input policyengine-core reads for ``name`` at ``period`` + in ``population``'s simulation, or None where no supplied input reaches + the period. Raises ``StoredDrawError`` where core would read anything + else, naming the variable, period, branch and tier. + + Changes nothing in storage and returns no array. + """ + schema = REGISTERED_DRAWS.get(name) + if schema is None: + raise StoredDrawError(f"{name} is not a registered stored draw.") + simulation = population.simulation + holder = population.get_holder(name) + period = periods.period(period) + _check_request(holder, name, schema, period) + _check_definition(simulation, holder, name, schema, period) + branch_name = simulation.branch_name + if holder._disk_storage is not None: + raise StoredDrawError( + f"The holder of {name} keeps values on disk, where " + f"policyengine-core cannot list inputs stored under a branch name " + f"with an underscore ({CORE_DISK_KEY_ISSUE}). Read {name} from a " + "simulation without memory_config." + ) + storing = holder._branch_storing(period, branch_name) + if storing is not None: + memory = holder._memory_storage + if not memory.is_derived(period, storing): + source = _supplied_source(simulation, holder, name, period, storing) + if source is None: + raise StoredDrawError( + f"policyengine-core would read " + f"{_describe(name, period, storing, 'memory')}, a value " + "set as an input that is not a supplied one." + ) + return source + cached = memory.get(period, storing) + source = _carry_source(simulation, holder, name, schema, period) + if source is not None: + if _same_array(cached, _source_array(holder, source)): + return source + raise StoredDrawError( + f"policyengine-core would read the value cached for " + f"{_describe(name, period, storing, 'memory')}, which differs " + f"from its supplied source " + f"{_describe(name, source.period, source.branch_name)}." + ) + if _same_array(cached, holder.default_array()): + return None + raise StoredDrawError( + f"policyengine-core would read the value cached for " + f"{_describe(name, period, storing, 'memory')}, but no supplied " + "input reaches the period and the value is not the default." + ) + source = _carry_source(simulation, holder, name, schema, period) + if source is None: + return None + if not simulation.tax_benefit_system.auto_carry_over_input_variables: + raise StoredDrawError( + f"The tax-benefit system does not carry inputs over, so " + f"policyengine-core would return the default for {name} in " + f"{period} instead of the input supplied for {source.period}." + ) + return source + + +def check_draw_answer(simulation, name: str, period, source: InputSource, raw): + """Raise unless ``raw``, core's answer for a read ``input_source`` traced + to ``source``, is that source's array, cast to the holder's type, and the + source is unchanged.""" + population = simulation.get_variable_population(name) + period = periods.period(period) + after = input_source(population, name, period) + if after != source: + raise StoredDrawError( + f"The supplied source of {name} for {period} changed while " + f"policyengine-core read it: {source} before, {after} after." + ) + holder = population.get_holder(name) + if not _same_array(raw, _source_array(holder, source)): + raise StoredDrawError( + f"policyengine-core returned a value for {name} in {period} on " + f"branch {simulation.branch_name!r} that is not its supplied " + f"source {_describe(name, source.period, source.branch_name, source.tier)}." + ) + + +def has_derived_values(holder) -> bool: + """Whether ``holder`` stores any value the simulation calculated.""" + for storage in (holder._memory_storage, holder._disk_storage): + if storage is None: + continue + for branch_name, stored_period in storage.get_known_branch_periods(): + if storage.is_derived(stored_period, branch_name): + return True + return False + + +def check_new_draw_input(simulation, name: str) -> None: + """Refuse a new input for a registered draw while a calculation runs, or + while the holder stores a value calculated from an earlier state.""" + if calculation_in_flight(): + raise StoredDrawError( + f"{name} cannot be set while a calculation is running: a stored " + "draw holds what the dataset supplies, and a calculation reads it." + ) + holder = simulation.get_holder(name) + if has_derived_values(holder): + raise StoredDrawError( + f"{name} cannot be set after the simulation has calculated a " + "value from it. Set the draw before reading it, or build a new " + "simulation." + ) + + +def check_option_request(simulation, name: str, method: str) -> None: + """Refuse a sum, twelfth or output callback of a registered draw.""" + raise StoredDrawError( + f"{name} is a stored draw, read only for its own period as stored, " + f"so {method} refuses it." + ) + + +# Stored outcome: whether a person takes up a student loan as a new borrower. +STUDENT_LOAN_TAKE_UP = "would_claim_student_loan" + + +def check_boolean_input(name: str, values) -> None: + """Raise unless every value is a boolean or the integer 0 or 1. + + policyengine-core skips its missing-value check for boolean variables and + casts every input to the variable's type, so a missing value would load + as True. Check the raw values first. + """ + array = np.asarray(values, dtype=object if _is_extension(values) else None) + flat = array.reshape(-1) + if array.dtype == bool: + return + if array.dtype.kind in "iu": + bad = ~np.isin(flat, (0, 1)) + else: + bad = np.array( + [ + not ( + isinstance(value, (bool, np.bool_)) + or ( + isinstance(value, (int, np.integer)) + and not isinstance(value, (bool, np.bool_)) + and value in (0, 1) + ) + ) + for value in flat + ], + dtype=bool, + ) + if bad.any(): + examples = ", ".join(repr(value) for value in flat[bad][:5]) + raise StoredDrawError( + f"{name} takes booleans or the integers 0 and 1, with nothing " + f"missing; {int(bad.sum())} value(s) are not ({examples})." + ) + + +def _is_extension(values) -> bool: + """A pandas extension array (nullable boolean, say), whose missing + values numpy would otherwise turn into something else.""" + dtype = getattr(values, "dtype", None) + return dtype is not None and not isinstance(dtype, np.dtype) From 0f9e471d9600202396f07041e1a90b194aa2e734 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 5 Nov 2026 07:02:51 +0000 Subject: [PATCH 05/11] Read stored draws in the UC modifier, student loan uprating and PIP scenario Each consumer asks input_source whether a stored input reaches the year. The UC rebalancing modifier compares uc_new_claimant_health_element_draw, promoted to float64, with the new-claimant share; the PIP scenario does the same with pip_phase_in_draw and the phase-in share; student loan uprating uses the stored would_claim_student_loan decision after checking its raw values. Without the column, each runs its existing seeded draw, inline. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LwQq6HUE9hj2BPEjEXLL68 --- policyengine_uk/data/economic_assumptions.py | 41 +- policyengine_uk/scenarios/pip_reform.py | 25 +- policyengine_uk/scenarios/uc_reform.py | 46 +- .../tests/test_stored_draw_consumers.py | 470 ++++++++++++++++++ 4 files changed, 561 insertions(+), 21 deletions(-) create mode 100644 policyengine_uk/tests/test_stored_draw_consumers.py diff --git a/policyengine_uk/data/economic_assumptions.py b/policyengine_uk/data/economic_assumptions.py index fd1f7c16fa..1db68dc26f 100644 --- a/policyengine_uk/data/economic_assumptions.py +++ b/policyengine_uk/data/economic_assumptions.py @@ -31,6 +31,10 @@ import logging from policyengine_uk.utils.stochastic import splitmix64_uniform +from policyengine_uk.utils.stored_draws import ( + STUDENT_LOAN_TAKE_UP, + check_boolean_input, +) # Approximate take-up rate for assigning loans to tertiary-educated NONE people. # This represents P(has loan AND earning above threshold | tertiary educated). @@ -824,6 +828,20 @@ def uprate_rent( return current_year +def stored_student_loan_take_up(person: pd.DataFrame) -> np.ndarray | None: + """The dataset's stored ``would_claim_student_loan`` for each row of + ``person``, or None where the dataset does not supply the column. + + Raises if a value is missing or is not a boolean or the integer 0 or 1: + cast to a boolean, a missing value would read as True. + """ + if STUDENT_LOAN_TAKE_UP not in person.columns: + return None + values = person[STUDENT_LOAN_TAKE_UP] + check_boolean_input(STUDENT_LOAN_TAKE_UP, values) + return np.asarray(values, dtype=bool) + + def uprate_student_loan_plans( current_year: UKSingleYearDataset, previous_year: UKSingleYearDataset, @@ -925,13 +943,20 @@ def uprate_student_loan_plans( new_plan = base_plan.copy() repayments = person["student_loan_repayments"].values.copy() - # One draw per person, the same in every year. - draws = splitmix64_uniform(person["person_id"].values, salt=_STUDENT_LOAN_DRAW_SALT) + # Whether each person takes up a loan as a new borrower: the dataset's + # stored decision where it supplies one, else one draw per person, the + # same in every year, below the take-up rate. + would_claim = stored_student_loan_take_up(person) + if would_claim is None: + draws = splitmix64_uniform( + person["person_id"].values, salt=_STUDENT_LOAN_DRAW_SALT + ) + would_claim = draws < _GRADUATE_LOAN_TAKE_UP # Helper to assign plans to eligible people - def assign_with_probability(mask, prob=_GRADUATE_LOAN_TAKE_UP): - """Assign the funding-nation plan to sampled new borrowers.""" - sampled = mask & (draws < prob) & (new_borrower_plan != "NONE") + def assign_new_borrowers(mask): + """Assign the funding-nation plan to new borrowers who take it up.""" + sampled = mask & would_claim & (new_borrower_plan != "NONE") new_plan[sampled] = new_borrower_plan[sampled] repayments[sampled] = 0.0 @@ -962,7 +987,7 @@ def assign_with_probability(mask, prob=_GRADUATE_LOAN_TAKE_UP): & (uni_start_year < _PLAN2_FIRST_UNI_START) & ~plan1_written_off ) - assign_with_probability(plan1_new_cohort) + assign_new_borrowers(plan1_new_cohort) # === Step 3: Add holders in the extended 2012-2022 age bands === # In the base year Plan 2 holders are at most the age of the 2012 intake @@ -977,7 +1002,7 @@ def assign_with_probability(mask, prob=_GRADUATE_LOAN_TAKE_UP): & (uni_start_year >= _PLAN2_FIRST_UNI_START) & (uni_start_year < _PLAN5_FIRST_UNI_START) ) - assign_with_probability(plan2_new_cohort) + assign_new_borrowers(plan2_new_cohort) # === Step 4: Add Plan 5 holders (new plan from 2023) === # Plan 5 is England's, and did not exist in the base year. The other @@ -993,7 +1018,7 @@ def assign_with_probability(mask, prob=_GRADUATE_LOAN_TAKE_UP): & (uni_start_year >= _PLAN5_FIRST_UNI_START) & (age >= 21) ) - assign_with_probability(plan5_eligible) + assign_new_borrowers(plan5_eligible) # Write back to the person table (without the merged region column). person_out = current_year.person.copy() diff --git a/policyengine_uk/scenarios/pip_reform.py b/policyengine_uk/scenarios/pip_reform.py index 8b632911ec..d8aada675c 100644 --- a/policyengine_uk/scenarios/pip_reform.py +++ b/policyengine_uk/scenarios/pip_reform.py @@ -1,20 +1,39 @@ from policyengine_uk.model_api import Scenario from policyengine_uk import Simulation +from policyengine_uk.utils.stored_draws import input_source import numpy as np +PHASE_IN_DRAW = "pip_phase_in_draw" + + +def phased_out_by_draw(draw, share) -> np.ndarray: + """People whose draw falls below the year's phase-out share, compared in + float64: a stored draw is a float32 value, and comparing it with a Python + float in float32 would round the share first.""" + return np.asarray(draw, dtype=np.float64) < share + def modify_simulation(sim: Simulation): - np.random.seed(42) # For reproducibility - pip_seed = np.random.random(len(sim.calculate("person_id"))) + # The seeded stream, drawn only for a year that no stored draw reaches, + # and then once: each person takes the value at their row. + pip_seed = None + draw_population = sim.get_variable_population(PHASE_IN_DRAW) start_year = 2025 end_year = 2029 for year in range(start_year, end_year + 1): + if input_source(draw_population, PHASE_IN_DRAW, year) is not None: + draw = sim.calculate(PHASE_IN_DRAW, year) + else: + if pip_seed is None: + np.random.seed(42) # For reproducibility + pip_seed = np.random.random(len(sim.calculate("person_id"))) + draw = pip_seed # Copy: calculate returns the simulation's cached array. current_pip = np.array(sim.calculate("pip", year)) percent_along_phase_in = (year - start_year) / (end_year - start_year) - current_pip[pip_seed < 0.25 * percent_along_phase_in] = 0 + current_pip[phased_out_by_draw(draw, 0.25 * percent_along_phase_in)] = 0 sim.set_input("pip", year, current_pip) return sim diff --git a/policyengine_uk/scenarios/uc_reform.py b/policyengine_uk/scenarios/uc_reform.py index eaa03ea02d..a9790813ae 100644 --- a/policyengine_uk/scenarios/uc_reform.py +++ b/policyengine_uk/scenarios/uc_reform.py @@ -3,6 +3,7 @@ from policyengine_uk.utils.data_source import built_from_data from policyengine_uk.data.economic_assumptions import PROJECTION_END_YEAR from policyengine_uk.utils.supplied_inputs import set_model_value, supplied_input +from policyengine_uk.utils.stored_draws import input_source import numpy as np import warnings from weakref import WeakKeyDictionary @@ -131,6 +132,7 @@ def _protected_existing_health_element_monthly( NEW_CLAIMANT_STATUS = "uc_receives_new_claimant_health_element" +NEW_CLAIMANT_DRAW = "uc_new_claimant_health_element_draw" LCWRA_AMOUNT = "gov.dwp.universal_credit.elements.disabled.amount" @@ -192,15 +194,15 @@ def warn_shadowed_lcwra_reform(sim) -> None: def _new_claimant_health_element_status( - sim: Microsimulation, year: int, drawn: np.ndarray, built_from_data: bool + sim: Microsimulation, year: int, drawn, built_from_data: bool ) -> np.ndarray: """Which benefit units get the new-claimant LCWRA rate in a year. A value the situation or dataset supplies is used as given. Data without - one takes the seeded draw, recorded as an input. A household situation - without one gets the protected amount (the variable's default): a draw - there would make a household's award depend on its position among the - other benefit units in the simulation. + one takes the draw (``drawn``, each benefit unit's decision), recorded + as an input. A household situation without one gets the protected amount + (the variable's default): a draw there would make a household's award + depend on its position among the other benefit units in the simulation. """ # sim.get_array reads the simulation's own branch (and the branches it # inherits from); the holder's get_array reads only the default branch. @@ -210,11 +212,20 @@ def _new_claimant_health_element_status( if built_from_data: status = np.asarray(drawn, dtype=bool) else: - status = np.zeros(len(drawn), dtype=bool) + status = np.zeros( + sim.get_variable_population(NEW_CLAIMANT_STATUS).count, dtype=bool + ) sim.set_input(NEW_CLAIMANT_STATUS, year, status) return status +def new_claimants_by_draw(draw, share) -> np.ndarray: + """Benefit units whose draw falls below the new-claimant share, compared + in float64: a stored draw is a float32 value, and comparing it with a + Python float in float32 would round the share first.""" + return np.asarray(draw, dtype=np.float64) < share + + def _lcwra_amount_is_overridden(sim, year: int) -> bool: """An explicit LCWRA amount keeps precedence over the inactive path.""" from policyengine_uk.utils.dated_reforms import legacy_override_bounds @@ -308,9 +319,9 @@ def add_universal_credit_reform(sim: Microsimulation): _add_inactive_health_element_formula(sim) rebalancing = sim.tax_benefit_system.parameters.gov.dwp.universal_credit.rebalancing - generator = np.random.default_rng(43) - - uc_seed = generator.random(len(sim.calculate("benunit_id"))) + # The seeded stream, drawn only for a year whose dataset supplies no + # stored draw, and then once: each unit takes its position's value. + uc_seed = None # rebalancing.new_claimant_share: WPI Economics' estimates for 2026-27 to # 2029-30, held after that (an assumption; see the parameter). new_claimant_share = rebalancing.new_claimant_share @@ -327,8 +338,23 @@ def add_universal_credit_reform(sim: Microsimulation): # Recalculate through the formula, preserving user-supplied awards. sim.delete_arrays("uc_LCWRA_element", year) continue + # Data without a status for the year takes the dataset's stored draw + # where one reaches the year, else the seeded stream. + drawn = None + if is_data and sim.get_array(NEW_CLAIMANT_STATUS, year) is None: + share = new_claimant_share(year) + draw_population = sim.get_variable_population(NEW_CLAIMANT_DRAW) + if input_source(draw_population, NEW_CLAIMANT_DRAW, year) is not None: + drawn = new_claimants_by_draw( + sim.calculate(NEW_CLAIMANT_DRAW, year), share + ) + else: + if uc_seed is None: + generator = np.random.default_rng(43) + uc_seed = generator.random(len(sim.calculate("benunit_id"))) + drawn = new_claimants_by_draw(uc_seed, share) is_post_2025_claimant = _new_claimant_health_element_status( - sim, year, uc_seed < new_claimant_share(year), is_data + sim, year, drawn, is_data ) # Copy: calculate returns the simulation's cached array, and the writes # below must reach the cache only through set_input. diff --git a/policyengine_uk/tests/test_stored_draw_consumers.py b/policyengine_uk/tests/test_stored_draw_consumers.py new file mode 100644 index 0000000000..e78bc13ae8 --- /dev/null +++ b/policyengine_uk/tests/test_stored_draw_consumers.py @@ -0,0 +1,470 @@ +"""The UC rebalancing modifier, student loan uprating and the PIP scenario +read the dataset's stored draws when it supplies them, and run their own +seeded draws when it does not. + +Invariants: + +1. Output-neutral without the columns: each consumer gives what its own draw + gives (the seed-43 stream, the person-id hash, the seed-42 stream). +2. With the columns, each decision is the stored value's: the draw compared + with the rate in float64, or the stored student loan decision. +3. Precision: a float32 stored draw compared with a share just above it keeps + its legacy side (3/8 - 2^-28 stored as 3/8 below 3/8 + 2^-28), and the + reference encoder's value for 0.0625 - 2^-30 stays below 2026's PIP + threshold where plain float32 rounding would not. +4. Reformability: new-claimant status is monotone in the share, and PIP + phase-out in the phase-in share; student loan take-up has no rate but the + module constant, and the stored decision is the same in every year. +5. Precedence: a status supplied for the year wins over the draw; one + supplied only earlier leaves the later year to the draw; a repeated + modifier, a reform after construction and an export replayed under + another share keep the statuses, as without stored draws. +6. Validation: would_claim_student_loan takes booleans or the integers 0 and + 1 only, with nothing missing, on every data route and in set_input. +7. Determinism: applying the PIP scenario with stored draws for every year + leaves numpy's global random state as it was. +""" + +import ast +import inspect + +import numpy as np +import pandas as pd +import pytest +from hypothesis import given, settings +from hypothesis import strategies as st + +import policyengine_uk.data.economic_assumptions as economic_assumptions +from policyengine_uk import Simulation +from policyengine_uk.data import UKSingleYearDataset +from policyengine_uk.data.economic_assumptions import ( + extend_single_year_dataset, + stored_student_loan_take_up, + uprate_student_loan_plans, +) +from policyengine_uk.model_api import Scenario +from policyengine_uk.scenarios.pip_reform import ( + modify_simulation, + phased_out_by_draw, + reform_pip_phase_in, +) +from policyengine_uk.scenarios.uc_reform import ( + REBALANCING_YEARS, + add_universal_credit_reform, + new_claimants_by_draw, +) +from policyengine_uk.system import system +from policyengine_uk.tests.stored_draw_fixtures import ( + PIP_THRESHOLDS, + SHARE_GRID, + encode, + one_person_units_frame, +) +from policyengine_uk.utils.stochastic import splitmix64_uniform +from policyengine_uk.utils.stored_draws import StoredDrawError + +STATUS = "uc_receives_new_claimant_health_element" +UC_DRAW = "uc_new_claimant_health_element_draw" +PIP_DRAW = "pip_phase_in_draw" +TAKE_UP = "would_claim_student_loan" +SHARE = "gov.dwp.universal_credit.rebalancing.new_claimant_share" +LEGACY = 3 / 8 - 2.0**-28 +STORED = np.float32(3 / 8) +ABOVE = 3 / 8 + 2.0**-28 +PROPERTY = settings(max_examples=200, deadline=None, derandomize=True) + + +def uc_frame(count, draws=None, status=None, years=(2025,)): + columns = { + "uc_limited_capability_for_WRA": {y: np.ones(count, bool) for y in years} + } + if draws is not None: + columns[UC_DRAW] = {years[0]: draws} + if status is not None: + columns[STATUS] = status + return one_person_units_frame(count=count, years=years, columns=columns) + + +def statuses(sim): + return { + year: np.asarray(sim.calculate(STATUS, year), dtype=bool) + for year in REBALANCING_YEARS + } + + +def shares(sim): + parameter = sim.tax_benefit_system.parameters.get_child(SHARE) + return {year: parameter(str(year)) for year in REBALANCING_YEARS} + + +# --- Universal Credit rebalancing ------------------------------------------ + + +@pytest.fixture(scope="module") +def uc_without_column(): + return Simulation(dataset=uc_frame(200)) + + +def test_uc_without_the_column_keeps_the_seed_43_stream(uc_without_column): + sim = uc_without_column + stream = np.random.default_rng(43).random(200) + share = shares(sim) + for year, status in statuses(sim).items(): + np.testing.assert_array_equal(status, stream < share[year]) + + +def test_uc_with_the_column_compares_the_stored_draw(): + stream = np.random.default_rng(43).random(200) + stored, feasible, _ = encode(stream, SHARE_GRID) + assert feasible.all() + with_column = Simulation(dataset=uc_frame(200, draws=stored)) + without = Simulation(dataset=uc_frame(200)) + share = shares(with_column) + got, legacy = statuses(with_column), statuses(without) + for year in REBALANCING_YEARS: + np.testing.assert_array_equal( + got[year], stored.astype(np.float64) < share[year] + ) + # The encoder keeps the legacy decision at every share on the grid. + np.testing.assert_array_equal(got[year], legacy[year]) + np.testing.assert_array_equal( + np.asarray(with_column.calculate("uc_LCWRA_element", year)), + np.asarray(without.calculate("uc_LCWRA_element", year)), + ) + + +def test_uc_promotes_the_stored_draw_before_comparing(): + assert LEGACY < ABOVE and float(STORED) == 3 / 8 + # numpy compares a float32 array with a Python float in float32. + assert not (np.array([STORED]) < ABOVE)[0] + assert new_claimants_by_draw(np.array([STORED]), ABOVE).tolist() == [True] + sim = Simulation( + dataset=uc_frame(1, draws=np.array([STORED])), + reform={SHARE: {"2026-01-01.2026-12-31": ABOVE}}, + ) + assert sim.calculate(STATUS, 2026).tolist() == [True] + + +@settings(PROPERTY) +@given( + draws=st.lists( + st.floats(0, 1, exclude_max=True, width=32), min_size=1, max_size=20 + ), + low=st.floats(0, 1), + high=st.floats(0, 1), +) +def test_new_claimant_status_is_monotone_in_the_share(draws, low, high): + low, high = sorted((low, high)) + draws = np.asarray(draws, dtype=np.float32) + below_low = new_claimants_by_draw(draws, low) + below_high = new_claimants_by_draw(draws, high) + assert not (below_low & ~below_high).any() + np.testing.assert_array_equal(below_low, draws.astype(np.float64) < low) + + +def test_uc_precedence_with_a_stored_draw(): + draws = np.array([0.05, 0.5, 0.95], dtype=np.float32) + supplied_2026 = np.array([False, True, False]) + sim = Simulation( + dataset=uc_frame( + 3, + draws=draws, + status={2026: supplied_2026}, + years=(2025, 2026), + ) + ) + share = shares(sim) + got = statuses(sim) + # Supplied for the year: wins over the draw. + np.testing.assert_array_equal(got[2026], supplied_2026) + # Supplied only for an earlier year: the later year takes the draw. + np.testing.assert_array_equal(got[2027], draws.astype(np.float64) < share[2027]) + # A repeated modifier reuses its own statuses. + add_universal_credit_reform(sim) + for year, status in statuses(sim).items(): + np.testing.assert_array_equal(status, got[year]) + # A reform after construction keeps them. + sim.apply_reform({SHARE: {"2027-01-01.2027-12-31": 1.0}}) + for year, status in statuses(sim).items(): + np.testing.assert_array_equal(status, got[year]) + + +def test_uc_export_replayed_under_another_share_keeps_the_statuses(): + draws = np.array([0.05, 0.5, 0.95], dtype=np.float32) + every_share = {SHARE: {"2026-01-01.2100-12-31": 1.0}} + no_share = {SHARE: {"2026-01-01.2100-12-31": 0.0}} + built = Simulation(dataset=uc_frame(3, draws=draws), reform=every_share) + exported = built.to_input_dataframe() + assert all(f"{STATUS}__{year}" in exported for year in REBALANCING_YEARS) + assert f"{UC_DRAW}__2025" in exported + replayed = Simulation(dataset=exported, reform=no_share) + for year, status in statuses(replayed).items(): + assert status.all(), year + + +def test_uc_situation_keeps_the_protected_amount_with_a_draw(): + situation = { + "people": {"adult": {"age": {2026: 30}}}, + "benunits": { + "benunit": {"members": ["adult"], UC_DRAW: {2026: 0.0}}, + }, + "households": {"household": {"members": ["adult"]}}, + } + sim = Simulation(situation=situation) + assert sim.calculate(STATUS, 2026).tolist() == [False] + + +# --- PIP phase-in ----------------------------------------------------------- + + +def pip_frame(count, draws=None): + years = tuple(range(2025, 2030)) + columns = {"pip": {year: np.full(count, 1_000.0) for year in years}} + if draws is not None: + columns[PIP_DRAW] = {2025: draws} + return one_person_units_frame(count=count, years=years, columns=columns) + + +def pip_by_year(sim): + return {year: np.asarray(sim.calculate("pip", year)) for year in range(2025, 2030)} + + +def test_pip_without_the_column_keeps_the_seed_42_stream(): + sim = Simulation(dataset=pip_frame(40), scenario=reform_pip_phase_in) + np.random.seed(42) + stream = np.random.random(40) + for year, pip in pip_by_year(sim).items(): + threshold = 0.25 * (year - 2025) / 4 + np.testing.assert_array_equal(pip == 0, stream < threshold) + + +def test_pip_with_the_column_matches_the_legacy_decisions_and_leaves_numpy_alone(): + np.random.seed(42) + stream = np.random.random(40) + stored, feasible, _ = encode(stream, PIP_THRESHOLDS) + assert feasible.all() + with_column = Simulation(dataset=pip_frame(40, draws=stored)) + without = Simulation(dataset=pip_frame(40)) + np.random.seed(7) + state = np.random.get_state() + reform_pip_phase_in.simulation_modifier(with_column) + after = np.random.get_state() + assert state[0] == after[0] and (state[1] == after[1]).all() + assert state[2:] == after[2:] + modify_simulation(without) + got, legacy = pip_by_year(with_column), pip_by_year(without) + for year in got: + np.testing.assert_array_equal(got[year], legacy[year]) + + +def test_pip_keeps_the_legacy_side_of_a_threshold_after_encoding(): + legacy = 0.0625 - 2.0**-30 + assert float(np.float32(legacy)) == 0.0625 + stored, feasible, _ = encode([legacy], PIP_THRESHOLDS) + assert feasible[0] and float(stored[0]) == 0.0625 - 2.0**-28 + sim = Simulation(dataset=pip_frame(1, draws=stored)) + modify_simulation(sim) + assert sim.calculate("pip", 2025).tolist() == [1_000] + assert sim.calculate("pip", 2026).tolist() == [0] + # Plain float32 rounding would have moved this person past 2026's share. + rounded = Simulation(dataset=pip_frame(1, draws=np.float32([legacy]))) + modify_simulation(rounded) + assert rounded.calculate("pip", 2026).tolist() == [1_000] + + +def test_pip_promotes_the_stored_draw_before_comparing(): + assert phased_out_by_draw(np.array([STORED]), ABOVE).tolist() == [True] + assert not (np.array([STORED]) < ABOVE)[0] + + +@settings(PROPERTY) +@given( + draws=st.lists( + st.floats(0, 1, exclude_max=True, width=32), min_size=1, max_size=20 + ), +) +def test_pip_phase_out_is_monotone_in_the_phase(draws): + draws = np.asarray(draws, dtype=np.float32) + previous = np.zeros(len(draws), dtype=bool) + for threshold in PIP_THRESHOLDS: + out = phased_out_by_draw(draws, threshold) + assert not (previous & ~out).any() + previous = out + + +def test_pip_situation_uses_a_stored_draw(): + situation = { + "people": { + "adult": { + "age": {2025: 30}, + "pip": {year: 1_000 for year in range(2025, 2030)}, + PIP_DRAW: {2025: 0.1}, + } + }, + "benunits": {"benunit": {"members": ["adult"]}}, + "households": {"household": {"members": ["adult"]}}, + } + sim = Simulation(situation=situation, scenario=reform_pip_phase_in) + assert [float(sim.calculate("pip", year)[0]) for year in range(2025, 2030)] == [ + 1_000, + 1_000, + 0, + 0, + 0, + ] + + +# --- Student loans ---------------------------------------------------------- + + +def loan_dataset(base_year=2024, take_up=None, ages=range(18, 51), per_age=20): + ages = np.repeat(list(ages), per_age) + ids = np.arange(1, len(ages) + 1) + person = pd.DataFrame( + { + "person_id": ids, + "person_benunit_id": ids, + "person_household_id": ids, + "age": ages, + "student_loan_plan": "NONE", + "highest_education": "TERTIARY", + "student_loan_repayments": 0.0, + } + ) + if take_up is not None: + person[TAKE_UP] = take_up(ids) + household = pd.DataFrame( + { + "household_id": ids, + "household_weight": 1.0, + "region": "LONDON", + "tenure_type": "OWNED_OUTRIGHT", + "rent": 0.0, + "council_tax": 0.0, + } + ) + return UKSingleYearDataset( + person=person, + benunit=pd.DataFrame({"benunit_id": ids}), + household=household, + fiscal_year=base_year, + ) + + +def plans(dataset): + projected = extend_single_year_dataset(dataset, system.parameters, end_year=2029) + return { + year: projected.datasets[year].person["student_loan_plan"].to_numpy() + for year in projected.datasets + } + + +def legacy_take_up(ids): + return splitmix64_uniform(np.asarray(ids), salt=3) < 0.4 + + +def test_the_stored_take_up_decision_is_todays_decision(): + without = plans(loan_dataset()) + stored = plans(loan_dataset(take_up=legacy_take_up)) + for year in without: + np.testing.assert_array_equal(stored[year], without[year]) + assert any((without[year] != "NONE").any() for year in without) + + +def test_the_stored_take_up_decision_replaces_the_hash(): + everyone = plans(loan_dataset(take_up=lambda ids: np.ones(len(ids), bool))) + nobody = plans(loan_dataset(take_up=lambda ids: np.zeros(len(ids), bool))) + base = loan_dataset() + for year in everyone: + if year == 2024: + continue + assert (nobody[year] == "NONE").all() + assert (everyone[year] != "NONE").sum() >= (plans(base)[year] != "NONE").sum() + + +def test_a_person_takes_up_in_every_projected_year_or_none(): + stored = loan_dataset(take_up=legacy_take_up) + took_up = stored.person[TAKE_UP].to_numpy() + for year, plan in plans(stored).items(): + assert ((plan != "NONE") <= took_up).all(), year + + +@pytest.mark.parametrize( + "values", + [ + [True, np.nan, False], + [1.0, 0.0, 1.0], + [2, 0, 1], + pd.array([True, None, False], dtype="boolean"), + ["True", "False", "True"], + ], + ids=["nan", "floats", "two", "nullable_missing", "strings"], +) +def test_take_up_values_are_checked_before_any_cast(values): + person = pd.DataFrame({TAKE_UP: values}) + with pytest.raises(StoredDrawError, match=TAKE_UP): + stored_student_loan_take_up(person) + sim = Simulation(dataset=one_person_units_frame(count=3)) + with pytest.raises(StoredDrawError, match=TAKE_UP): + sim.set_input(TAKE_UP, 2025, values) + assert sim.get_holder(TAKE_UP).get_known_periods() == [] + + +@pytest.mark.parametrize( + "values", + [[True, False, True], [1, 0, 1], np.array([1, 0, 1], dtype=np.int8)], +) +def test_take_up_accepts_booleans_and_zero_or_one(values): + np.testing.assert_array_equal( + stored_student_loan_take_up(pd.DataFrame({TAKE_UP: values})), + [True, False, True], + ) + sim = Simulation(dataset=one_person_units_frame(count=3)) + sim.set_input(TAKE_UP, 2025, values) + assert sim.calculate(TAKE_UP, 2025).tolist() == [True, False, True] + + +def test_a_dataset_with_a_missing_take_up_value_is_refused_on_each_route(): + dataset = loan_dataset(take_up=lambda ids: np.where(ids == 1, np.nan, 1.0)) + with pytest.raises(StoredDrawError, match=TAKE_UP): + uprate_student_loan_plans(dataset.copy(), dataset, system.parameters) + with pytest.raises(StoredDrawError, match=TAKE_UP): + Simulation(dataset=dataset) + frame = one_person_units_frame(count=2, columns={TAKE_UP: [True, None]}) + with pytest.raises(StoredDrawError, match=TAKE_UP): + Simulation(dataset=frame) + + +def test_no_take_up_call_passes_a_rate(): + """The stored decision is today's only at the module constant; a call + passing another rate would need the stored draw instead.""" + assert economic_assumptions._GRADUATE_LOAN_TAKE_UP == 0.4 + tree = ast.parse(inspect.getsource(uprate_student_loan_plans)) + calls = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "assign_new_borrowers" + ] + assert len(calls) == 3 + assert all(len(call.args) == 1 and not call.keywords for call in calls) + helper = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.FunctionDef) and node.name == "assign_new_borrowers" + ) + assert [argument.arg for argument in helper.args.args] == ["mask"] + rates = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Name) and node.id == "_GRADUATE_LOAN_TAKE_UP" + ] + assert len(rates) == 1 + module = ast.parse(inspect.getsource(economic_assumptions)) + readers = [ + node + for node in ast.walk(module) + if isinstance(node, ast.Name) and node.id == "_GRADUATE_LOAN_TAKE_UP" + ] + assert len(readers) == 2 # the assignment and the fallback's comparison From 276d7caf0c00dd015f80bee128ce9a2de3490e6d Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 5 Nov 2026 07:10:49 +0000 Subject: [PATCH 06/11] Check identifiers and supplied periods on the core Dataset route; gate core versions The core Dataset route refuses a dataset carrying a stored column whose person, benefit unit or household identifiers do not survive the int32 cast, and one whose stored columns or exempting outcomes are supplied for a later period but not the first, or whose draws change between periods, before it exports anything. A code-health test fails when uv.lock or the installed policyengine-core holds a version outside the qualified set, and another fails on any package read of a stored draw that skips calculate. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LwQq6HUE9hj2BPEjEXLL68 --- changelog.d/2874.added.md | 1 + policyengine_uk/data/core_route_checks.py | 154 ++++++++++ policyengine_uk/simulation.py | 29 +- .../code_health/test_core_version_gate.py | 81 ++++++ .../code_health/test_stored_draw_reads.py | 103 +++++++ .../tests/test_stored_draw_consumers.py | 2 + .../tests/test_stored_draw_core_route.py | 263 ++++++++++++++++++ .../tests/test_stored_draw_input_source.py | 84 ++++++ .../tests/test_supplied_input_tightening.py | 2 + policyengine_uk/utils/stored_draws.py | 2 + 10 files changed, 712 insertions(+), 9 deletions(-) create mode 100644 changelog.d/2874.added.md create mode 100644 policyengine_uk/data/core_route_checks.py create mode 100644 policyengine_uk/tests/code_health/test_core_version_gate.py create mode 100644 policyengine_uk/tests/code_health/test_stored_draw_reads.py create mode 100644 policyengine_uk/tests/test_stored_draw_core_route.py diff --git a/changelog.d/2874.added.md b/changelog.d/2874.added.md new file mode 100644 index 0000000000..f8e8d7d8d0 --- /dev/null +++ b/changelog.d/2874.added.md @@ -0,0 +1 @@ +Simulations read three new stored inputs, `uc_new_claimant_health_element_draw`, `pip_phase_in_draw` and `would_claim_student_loan`, in Universal Credit rebalancing, the PIP phase-in scenario and student loan uprating when a dataset supplies them, and keep their own seeded draws when it does not. diff --git a/policyengine_uk/data/core_route_checks.py b/policyengine_uk/data/core_route_checks.py new file mode 100644 index 0000000000..16977b5704 --- /dev/null +++ b/policyengine_uk/data/core_route_checks.py @@ -0,0 +1,154 @@ +"""Checks the policyengine-core ``Dataset`` route makes before it exports. + +``Simulation.build_from_dataset`` stores each variable of a core ``Dataset`` +through ``set_input``, which casts identifiers to int32, exports the +simulation at the dataset's first period with +``UKSingleYearDataset.from_simulation`` and rebuilds from that export. Two +things a stored draw or decision needs can be lost on the way: + +* Identities. The rebuild declares benefit units and households in sorted + order of the cast identifiers, and a stored decision keyed on a raw + identifier no longer matches one keyed on the cast identifier. So a dataset + carrying a stored column must give identifiers that survive the int32 cast, + for people, benefit units and households. +* Periods. Only the first period is exported. A column supplied only for + later periods would be exported as its default and then stored as a + supplied input. So a stored column, or an outcome that exempts a draw, that + the dataset supplies for any period must be supplied for the first; and a + draw's later periods must repeat the first, since only the first is kept. +""" + +import warnings + +import numpy as np +from policyengine_core.periods import period as period_ + +# Stored columns the data build writes: the UC deduction draws, the +# birthday position, the Warm Home Discount draw, the UC new-claimant draw, +# the student loan decision and the PIP phase-in draw. +STORED_COLUMNS = ( + "uc_deduction_random_draw", + "uc_deduction_type_random_draw", + "months_since_last_birthday", + "warm_home_discount_random_draw", + "uc_new_claimant_health_element_draw", + "would_claim_student_loan", + "pip_phase_in_draw", +) +# Draws, whose later periods must repeat the first. +DRAW_COLUMNS = ( + "uc_deduction_random_draw", + "uc_deduction_type_random_draw", + "warm_home_discount_random_draw", + "uc_new_claimant_health_element_draw", + "pip_phase_in_draw", + "attends_private_school_random_draw", +) +# Columns that must be supplied for the first period if supplied at all: the +# stored columns, the private school draw, and the outcomes that exempt a +# draw. +FIRST_PERIOD_COLUMNS = STORED_COLUMNS + ( + "attends_private_school_random_draw", + "uc_receives_new_claimant_health_element", + "attends_private_school", +) +# The identifier arrays of each entity: its own and its members' memberships. +ENTITY_IDENTIFIERS = { + "person": ("person_id",), + "benunit": ("benunit_id", "person_benunit_id"), + "household": ("household_id", "person_household_id"), +} + + +class CoreRouteError(ValueError): + """A core ``Dataset`` that the export and rebuild would change.""" + + +def _fails_int32_cast(ids) -> np.ndarray: + """Which identifiers change when cast to int32.""" + values = np.asarray(ids) + if values.dtype.kind == "O": + try: + values = np.array(values.tolist(), dtype=np.float64) + except (TypeError, ValueError): + return np.ones(values.shape, dtype=bool) + if values.dtype.kind not in "iub": + values = values.astype(np.float64) + finite = np.isfinite(values) + safe = np.where(finite, values, 0.0) + in_range = (safe >= np.iinfo(np.int32).min) & (safe <= np.iinfo(np.int32).max) + return ~(finite & in_range & (safe == np.trunc(safe))) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + return values.astype(np.int32).astype(values.dtype) != values + + +def check_identifiers(data: dict) -> None: + """Refuse a dataset carrying a stored column whose identifiers do not + survive the int32 cast, naming the entity and the identifiers.""" + carried = [column for column in STORED_COLUMNS if column in data] + if not carried: + return + for entity, columns in ENTITY_IDENTIFIERS.items(): + failing = [] + for column in columns: + for values in (data.get(column) or {}).values(): + values = np.asarray(values) + bad = _fails_int32_cast(values) + failing.extend(values[bad].reshape(-1).tolist()) + if failing: + unique = list(dict.fromkeys(failing)) + raise CoreRouteError( + f"This core Dataset carries {', '.join(carried)}, and " + f"{len(unique)} {entity} identifier(s) do not survive the " + f"int32 cast its route stores identifiers with (the first: " + f"{unique[:5]}). The rebuild would key or order the stored " + "values on other identifiers. Renumber the identifiers within " + "the int32 range, or load the data as a UKSingleYearDataset." + ) + + +def check_supplied_periods( + data: dict, system, first_time_period, group_orders: dict +) -> None: + """Refuse a dataset whose stored columns the first-period export would + change, naming the column and its periods.""" + first = str(period_(first_time_period)) + for column in FIRST_PERIOD_COLUMNS: + if column not in data or column not in system.variables: + continue + by_period = data[column] + if not len(by_period): + continue + supplied = {str(period_(key)): key for key in by_period} + if first not in supplied: + raise CoreRouteError( + f"This core Dataset supplies {column} for " + f"{', '.join(sorted(supplied))} but not for {first}, the only " + "period its route exports. The export would store the " + "variable's default as a supplied value." + ) + if column not in DRAW_COLUMNS or len(supplied) == 1: + continue + variable = system.variables[column] + orders = group_orders.get(variable.entity.key) + + def aligned(key): + values = np.asarray(by_period[key]) + if orders is not None and values.ndim > 0: + values = values[orders.get(key, orders[supplied[first]])] + return values.astype(variable.dtype) + + first_values = aligned(supplied[first]) + differing = [ + name + for name, key in supplied.items() + if name != first and not np.array_equal(aligned(key), first_values) + ] + if differing: + raise CoreRouteError( + f"This core Dataset supplies {column} for {', '.join(sorted(supplied))} " + f"with values in {', '.join(sorted(differing))} that differ " + f"from {first}'s. Its route keeps only {first}, so a draw " + "must not change between periods." + ) diff --git a/policyengine_uk/simulation.py b/policyengine_uk/simulation.py index 24c762a01a..c6b71a3f1e 100644 --- a/policyengine_uk/simulation.py +++ b/policyengine_uk/simulation.py @@ -45,6 +45,10 @@ uprating_input_fingerprint, ) from policyengine_uk.data.dataset_sources import materialize_gcs_dataset_url +from policyengine_uk.data.core_route_checks import ( + check_identifiers, + check_supplied_periods, +) from policyengine_uk.data.energy_price_basis import core_dataset_energy_price_basis from policyengine_uk.data.savings_income_projection import ( SAVINGS_INCOME_PROJECTION_KEY, @@ -1445,6 +1449,22 @@ def get_first_array(variable_name: str) -> np.ndarray: time_period_values = data[variable_name] return time_period_values[first_time_period] + # Populations use sorted group IDs; arrays in a Core dataset follow + # the order of that period's supplied identifier array instead. + group_orders = { + entity: { + period: np.argsort(ids) for period, ids in data[f"{entity}_id"].items() + } + for entity in ("benunit", "household") + } + # The export below keeps only the first period, and the rebuild + # re-declares groups by their int32 identifiers: refuse data whose + # stored columns either would change. + check_identifiers(data) + check_supplied_periods( + data, self.tax_benefit_system, first_time_period, group_orders + ) + # Build entity structure from IDs self.build_from_ids( *map( @@ -1466,15 +1486,6 @@ def get_first_array(variable_name: str) -> np.ndarray: ) self._savings_income_projection = savings_income_projection - # Populations use sorted group IDs; arrays in a Core dataset follow - # the order of that period's supplied identifier array instead. - group_orders = { - entity: { - period: np.argsort(ids) for period, ids in data[f"{entity}_id"].items() - } - for entity in ("benunit", "household") - } - # Load all variable values for variable in data: for time_period in data[variable]: diff --git a/policyengine_uk/tests/code_health/test_core_version_gate.py b/policyengine_uk/tests/code_health/test_core_version_gate.py new file mode 100644 index 0000000000..f75d113813 --- /dev/null +++ b/policyengine_uk/tests/code_health/test_core_version_gate.py @@ -0,0 +1,81 @@ +"""policyengine-uk runs only on a policyengine-core version it has qualified. + +The stored-draw check (policyengine_uk.utils.stored_draws) reads core's +private members and relies on how core stores, carries and caches inputs. +Tests qualify that behaviour for each exact core version in +QUALIFIED_CORE_VERSIONS. The resolver floor in pyproject.toml does not. + +Both checks fail on any other version, including a new core that a lock +refresh brings in. Before adding a version to QUALIFIED_CORE_VERSIONS, run +these on that exact version and record the result: + + uv run --frozen --with 'policyengine-core==' pytest \ + policyengine_uk/tests/test_stored_draw_input_source.py \ + policyengine_uk/tests/test_stored_draw_consumers.py \ + policyengine_uk/tests/test_stored_draw_core_route.py \ + policyengine_uk/tests/test_supplied_input_tightening.py -n 4 + +The checks run when the tests run, not when policyengine_uk is imported. +""" + +import importlib.metadata +import re +from pathlib import Path + +from policyengine_uk.utils.stored_draws import QUALIFIED_CORE_VERSIONS + +PACKAGE = Path(__file__).resolve().parents[2] +RERUN = ( + "Run the stored-draw qualification tests on that exact version (see this " + "file's docstring) before adding it to QUALIFIED_CORE_VERSIONS in " + "policyengine_uk/utils/stored_draws.py." +) + + +def locked_core_version(lock_text: str) -> str: + match = re.search( + r'^\[\[package\]\]\nname = "policyengine-core"\nversion = "([^"]+)"', + lock_text, + re.MULTILINE, + ) + assert match, "uv.lock names no policyengine-core package." + return match.group(1) + + +def find_lock() -> Path: + """uv.lock beside the source tree, or in the directory the tests run + from: a test of the built wheel runs from the checkout, whose lock + installed the environment.""" + for directory in (PACKAGE.parent, Path.cwd()): + lock = directory / "uv.lock" + if lock.is_file(): + return lock + raise AssertionError( + f"No uv.lock beside {PACKAGE} or in {Path.cwd()}: run the tests from " + "the checkout whose lock installed this environment." + ) + + +def test_the_locked_core_is_qualified(): + version = locked_core_version(find_lock().read_text()) + assert version in QUALIFIED_CORE_VERSIONS, ( + f"uv.lock holds policyengine-core {version}, which is not in " + f"{sorted(QUALIFIED_CORE_VERSIONS)}. {RERUN}" + ) + + +def test_the_installed_core_is_qualified(): + version = importlib.metadata.version("policyengine-core") + assert version in QUALIFIED_CORE_VERSIONS, ( + f"policyengine-core {version} is installed, which is not in " + f"{sorted(QUALIFIED_CORE_VERSIONS)}. {RERUN}" + ) + + +def test_the_lock_parser_reads_the_core_entry(): + lock = ( + '[[package]]\nname = "policyengine-core-extra"\nversion = "9.9.9"\n\n' + '[[package]]\nname = "policyengine-core"\nversion = "3.34.3"\n' + ) + assert locked_core_version(lock) == "3.34.3" + assert "3.34.3" not in QUALIFIED_CORE_VERSIONS diff --git a/policyengine_uk/tests/code_health/test_stored_draw_reads.py b/policyengine_uk/tests/code_health/test_stored_draw_reads.py new file mode 100644 index 0000000000..6296bc4893 --- /dev/null +++ b/policyengine_uk/tests/code_health/test_stored_draw_reads.py @@ -0,0 +1,103 @@ +"""pe-uk reads a stored draw only through Simulation.calculate. + +``Simulation.calculate`` checks each read of a registered draw against the +supplied input policyengine-core reads (``utils/stored_draws.input_source``). +A read that goes to storage instead (``get_array``, +``_get_array_from_storage``, a holder) skips that check. This scan fails on +any such read of a registered draw in the package, named by a string or by +a module constant that holds the name. ``utils/stored_draws.py`` is the check +itself and reads storage on purpose. +""" + +import ast +from pathlib import Path + +from policyengine_uk.utils.stored_draws import REGISTERED_DRAW_NAMES + +PACKAGE = Path(__file__).resolve().parents[2] +EXEMPT = {"utils/stored_draws.py"} +STORAGE_READS = { + "get_array", + "_get_array_from_storage", + "get_holder", + "get_stored_input", + "stored_supplied_value", + "supplied_input", + "get_known_periods", + "get_input_periods", +} + + +def _draw_constants(tree) -> set: + """Module-level names bound to a registered draw's name.""" + names = set() + for node in tree.body: + if ( + isinstance(node, ast.Assign) + and isinstance(node.value, ast.Constant) + and node.value.value in REGISTERED_DRAW_NAMES + ): + names.update(t.id for t in node.targets if isinstance(t, ast.Name)) + return names + + +def _names_a_draw(node, constants) -> bool: + return (isinstance(node, ast.Constant) and node.value in REGISTERED_DRAW_NAMES) or ( + isinstance(node, ast.Name) and node.id in constants + ) + + +def storage_reads(source: str) -> list: + tree = ast.parse(source) + constants = _draw_constants(tree) + found = [] + for node in ast.walk(tree): + if isinstance(node, ast.Call): + func = node.func + name = ( + func.attr + if isinstance(func, ast.Attribute) + else getattr(func, "id", None) + ) + arguments = list(node.args) + [keyword.value for keyword in node.keywords] + if name in STORAGE_READS and any( + _names_a_draw(argument, constants) for argument in arguments + ): + found.append((node.lineno, name)) + if ( + isinstance(node, ast.Subscript) + and isinstance(node.value, ast.Attribute) + and node.value.attr == "_holders" + and _names_a_draw(node.slice, constants) + ): + found.append((node.lineno, "_holders")) + return found + + +def test_no_package_module_reads_a_draw_from_storage(): + found = [] + for path in sorted(PACKAGE.rglob("*.py")): + relative = path.relative_to(PACKAGE).as_posix() + if relative.startswith("tests/") or relative in EXEMPT: + continue + for line, name in storage_reads(path.read_text()): + found.append(f"{relative}:{line} {name}") + assert not found, ( + "These read a stored draw from storage, skipping the check in " + "Simulation.calculate; read it with calculate instead:\n" + "\n".join(found) + ) + + +def test_the_scan_finds_direct_reads(): + source = ( + 'DRAW = "pip_phase_in_draw"\n' + "def read(sim, year):\n" + " sim.get_array(DRAW, year)\n" + ' sim.get_holder("uc_new_claimant_health_element_draw")\n' + ' population._holders["pip_phase_in_draw"]\n' + " holder._get_array_from_storage(DRAW, year)\n" + " sim.calculate(DRAW, year)\n" + ) + assert sorted(name for _, name in storage_reads(source)) == sorted( + ["get_array", "get_holder", "_holders", "_get_array_from_storage"] + ) diff --git a/policyengine_uk/tests/test_stored_draw_consumers.py b/policyengine_uk/tests/test_stored_draw_consumers.py index e78bc13ae8..817dac9a98 100644 --- a/policyengine_uk/tests/test_stored_draw_consumers.py +++ b/policyengine_uk/tests/test_stored_draw_consumers.py @@ -63,6 +63,8 @@ from policyengine_uk.utils.stochastic import splitmix64_uniform from policyengine_uk.utils.stored_draws import StoredDrawError +pytestmark = pytest.mark.usefixtures("cloned_uk_tax_benefit_system") + STATUS = "uc_receives_new_claimant_health_element" UC_DRAW = "uc_new_claimant_health_element_draw" PIP_DRAW = "pip_phase_in_draw" diff --git a/policyengine_uk/tests/test_stored_draw_core_route.py b/policyengine_uk/tests/test_stored_draw_core_route.py new file mode 100644 index 0000000000..c66f5f2427 --- /dev/null +++ b/policyengine_uk/tests/test_stored_draw_core_route.py @@ -0,0 +1,263 @@ +"""Stored columns on policyengine-core's ``Dataset`` route, and the +with/without differential on both loading routes. + +Invariants: + +1. Keys and order: the route refuses a dataset carrying a stored column + where any person, benefit unit or household identifier does not survive + the int32 cast, naming the entity; without a stored column such a dataset + loads as before. +2. Supplied periods: a stored column or exempting outcome supplied for any + period must be supplied for the first, and a draw's later periods must + repeat its first, or the route refuses the dataset before it exports + anything. An outcome may vary across periods; later periods drop, as for + any input. +3. Neutral by construction: with columns the reference encoder built from + today's draws, each loading route gives the same UC new-claimant + statuses, LCWRA amounts and PIP amounts, bit for bit, as without them. +4. The stored columns survive the route: the rebuilt simulation reads the + values the file supplied. +""" + +import h5py +import numpy as np +import pytest + +from policyengine_uk import Simulation +from policyengine_uk.data.core_route_checks import CoreRouteError +from policyengine_uk.scenarios.pip_reform import modify_simulation +from policyengine_uk.scenarios.uc_reform import REBALANCING_YEARS +from policyengine_uk.system import system +from policyengine_uk.tests.stored_draw_fixtures import ( + ID_COLUMNS, + PIP_THRESHOLDS, + SHARE_GRID, + encode, + one_person_units_frame, +) + +pytestmark = pytest.mark.usefixtures("cloned_uk_tax_benefit_system") + +STATUS = "uc_receives_new_claimant_health_element" +UC_DRAW = "uc_new_claimant_health_element_draw" +PIP_DRAW = "pip_phase_in_draw" +INT32_MAX = np.iinfo(np.int32).max + + +def add_households(file, year, count): + """Household columns the route's export and projection read.""" + region = system.variables["region"].possible_values + tenure = system.variables["tenure_type"].possible_values + file.create_dataset(f"household_weight/{year}", data=np.ones(count)) + file.create_dataset( + f"region/{year}", data=np.asarray(region.encode(np.full(count, "LONDON"))) + ) + file.create_dataset( + f"tenure_type/{year}", + data=np.asarray(tenure.encode(np.full(count, "OWNED_OUTRIGHT"))), + ) + file.create_dataset(f"rent/{year}", data=np.zeros(count)) + file.create_dataset(f"council_tax/{year}", data=np.zeros(count)) + + +def core_file(path, count=3, columns=None, ids=None, years=(2025,)): + """A core Dataset file: each variable's array under ``variable/year``, + one adult per benefit unit and household.""" + ids = {name: np.arange(1, count + 1) for name in ID_COLUMNS} | (ids or {}) + with h5py.File(path, "w") as file: + for year in years: + for name, values in ids.items(): + file.create_dataset(f"{name}/{year}", data=np.asarray(values)) + file.create_dataset(f"age/{year}", data=np.full(count, 30)) + add_households(file, year, count) + for name, by_year in (columns or {}).items(): + for year, values in by_year.items(): + file.create_dataset(f"{name}/{year}", data=np.asarray(values)) + return str(path) + + +# --- Identities ------------------------------------------------------------- + + +@pytest.mark.parametrize( + "entity, ids", + [ + ("person", {"person_id": [1, 2, INT32_MAX + 1]}), + ( + "benunit", + { + "benunit_id": [1, 2, INT32_MAX + 1], + "person_benunit_id": [1, 2, INT32_MAX + 1], + }, + ), + ( + "household", + { + "household_id": [1, 2, INT32_MAX + 1], + "person_household_id": [1, 2, INT32_MAX + 1], + }, + ), + ], +) +@pytest.mark.parametrize("column", [UC_DRAW, PIP_DRAW, "months_since_last_birthday"]) +def test_identifiers_beyond_int32_refuse_a_stored_column(entity, ids, column, tmp_path): + values = {UC_DRAW: [0.5] * 3, PIP_DRAW: [0.5] * 3}.get(column, [6.0] * 3) + path = core_file(tmp_path / "core.h5", columns={column: {2025: values}}, ids=ids) + with pytest.raises(CoreRouteError, match=f"{entity} identifier"): + Simulation(dataset=path) + + +def test_identifiers_beyond_int32_load_without_a_stored_column(tmp_path): + path = core_file( + tmp_path / "core.h5", + ids={"person_id": [1, 2, INT32_MAX + 1]}, + ) + with pytest.warns(Warning): + sim = Simulation(dataset=path) + assert len(sim.calculate("person_id", 2025)) == 3 + + +def test_benefit_unit_order_reversed_by_the_cast_is_refused(tmp_path): + # Raw [1, 2^31] sorts the other way once 2^31 wraps to -2^31. + ids = {"benunit_id": [1, 2**31], "person_benunit_id": [1, 2**31]} + path = core_file( + tmp_path / "core.h5", + count=2, + ids={ + **ids, + **{k: [1, 2] for k in ("person_id", "household_id", "person_household_id")}, + }, + columns={UC_DRAW: {2025: [0.1, 0.9]}}, + ) + with pytest.raises(CoreRouteError, match="benunit identifier"): + Simulation(dataset=path) + + +# --- Supplied periods ------------------------------------------------------- + + +@pytest.mark.parametrize( + "column, values", + [ + (PIP_DRAW, [0.5] * 3), + # Its default, 0.0, is a legal value: the export would store it. + ("warm_home_discount_random_draw", [0.0] * 3), + ("would_claim_student_loan", [True] * 3), + (STATUS, [True] * 3), + ("attends_private_school", [False] * 3), + ], +) +def test_a_column_supplied_only_for_a_later_period_is_refused(column, values, tmp_path): + path = core_file( + tmp_path / "core.h5", columns={column: {2026: values}}, years=(2025, 2026) + ) + with pytest.raises(CoreRouteError, match=f"{column} for 2026 but not for 2025"): + Simulation(dataset=path) + + +def test_a_draw_that_changes_between_periods_is_refused(tmp_path): + path = core_file( + tmp_path / "core.h5", + columns={PIP_DRAW: {2025: [0.1, 0.2, 0.3], 2026: [0.1, 0.2, 0.4]}}, + years=(2025, 2026), + ) + with pytest.raises(CoreRouteError, match="differ from 2025"): + Simulation(dataset=path) + + +def test_a_repeated_draw_and_a_changing_outcome_load(tmp_path): + path = core_file( + tmp_path / "core.h5", + columns={ + PIP_DRAW: {2025: [0.1, 0.2, 0.3], 2026: [0.1, 0.2, 0.3]}, + STATUS: {2025: [True, False, True], 2026: [False, False, True]}, + }, + years=(2025, 2026), + ) + sim = Simulation(dataset=path) + np.testing.assert_array_equal( + sim.calculate(PIP_DRAW, 2027), np.float32([0.1, 0.2, 0.3]) + ) + + +def test_a_group_draw_repeated_in_another_row_order_loads(tmp_path): + ids_2026 = {"benunit_id": [3, 1, 2]} + with h5py.File(tmp_path / "core.h5", "w") as file: + for year in (2025, 2026): + for name in ID_COLUMNS: + values = np.arange(1, 4) + if year == 2026 and name in ids_2026: + values = np.asarray(ids_2026[name]) + file.create_dataset(f"{name}/{year}", data=values) + file.create_dataset(f"age/{year}", data=np.full(3, 30)) + add_households(file, year, 3) + file.create_dataset(f"{UC_DRAW}/2025", data=np.float32([0.1, 0.2, 0.3])) + file.create_dataset(f"{UC_DRAW}/2026", data=np.float32([0.3, 0.1, 0.2])) + sim = Simulation(dataset=str(tmp_path / "core.h5")) + np.testing.assert_array_equal( + sim.calculate(UC_DRAW, 2025), np.float32([0.1, 0.2, 0.3]) + ) + + +# --- With/without differential on both routes -------------------------------- + + +def legacy_streams(count): + uc = np.random.default_rng(43).random(count) + np.random.seed(42) + pip = np.random.random(count) + return uc, pip + + +def outputs(sim): + modify_simulation(sim) + result = {} + for year in REBALANCING_YEARS: + result[(STATUS, year)] = np.asarray(sim.calculate(STATUS, year)) + result[("uc_LCWRA_element", year)] = np.asarray( + sim.calculate("uc_LCWRA_element", year) + ) + for year in range(2025, 2030): + result[("pip", year)] = np.asarray(sim.calculate("pip", year)) + return result + + +@pytest.mark.parametrize("route", ["frame", "core"]) +def test_with_and_without_the_columns_agree_bit_for_bit(route, tmp_path): + count = 120 + uc, pip = legacy_streams(count) + uc_stored, uc_ok, _ = encode(uc, SHARE_GRID) + pip_stored, pip_ok, _ = encode(pip, PIP_THRESHOLDS) + assert uc_ok.all() and pip_ok.all() + base = { + "uc_limited_capability_for_WRA": {2025: np.ones(count, bool)}, + "pip": {2025: np.full(count, 1_000.0)}, + } + stored = {UC_DRAW: {2025: uc_stored}, PIP_DRAW: {2025: pip_stored}} + + def build(columns, name): + if route == "frame": + return Simulation( + dataset=one_person_units_frame(count=count, columns=columns) + ) + return Simulation( + dataset=core_file(tmp_path / f"{name}.h5", count=count, columns=columns) + ) + + without = outputs(build(base, "without")) + with_columns_sim = build(base | stored, "with") + np.testing.assert_array_equal(with_columns_sim.calculate(UC_DRAW, 2025), uc_stored) + np.testing.assert_array_equal( + with_columns_sim.calculate(PIP_DRAW, 2025), pip_stored + ) + with_columns = outputs(with_columns_sim) + assert without.keys() == with_columns.keys() + for key in without: + assert without[key].dtype == with_columns[key].dtype, key + np.testing.assert_array_equal(with_columns[key], without[key], err_msg=str(key)) + # The decisions are the legacy streams' own. + for year in REBALANCING_YEARS: + share = with_columns_sim.tax_benefit_system.parameters.get_child( + "gov.dwp.universal_credit.rebalancing.new_claimant_share" + )(str(year)) + np.testing.assert_array_equal(without[(STATUS, year)], uc < share) diff --git a/policyengine_uk/tests/test_stored_draw_input_source.py b/policyengine_uk/tests/test_stored_draw_input_source.py index e9c2b98794..86e906c52e 100644 --- a/policyengine_uk/tests/test_stored_draw_input_source.py +++ b/policyengine_uk/tests/test_stored_draw_input_source.py @@ -49,6 +49,8 @@ input_source, ) +pytestmark = pytest.mark.usefixtures("cloned_uk_tax_benefit_system") + DRAW = "pip_phase_in_draw" GROUP_DRAW = "uc_new_claimant_health_element_draw" VALUES = np.array([0.25, 0.5, 0.75], dtype=np.float32) @@ -378,6 +380,7 @@ def test_redefinitions_raise(): {"formula": lambda person, period, parameters: 0.5}, {"adds": ["age"]}, {"value_type": int, "default_value": 1}, + {"requires_computation_after": "age"}, ): sim = built({2025: VALUES}) body = dict( @@ -394,6 +397,14 @@ def test_redefinitions_raise(): sim.calculate(DRAW, p) +def test_a_draw_removed_from_the_system_raises(): + sim = built({2025: VALUES}) + population = sim.get_variable_population(DRAW) + del sim.tax_benefit_system.variables[DRAW] + with pytest.raises(StoredDrawError, match="has no current definition"): + input_source(population, DRAW, 2025) + + def test_a_draw_moved_to_another_entity_raises(): sim = built({2025: VALUES}) sim.tax_benefit_system.update_variable( @@ -611,3 +622,76 @@ def test_situations_keep_cores_path(): np.testing.assert_array_equal( sim.calculate_divide(DRAW, "2025-06"), np.float32([1 / 12]) ) + + +def test_two_inputs_starting_on_one_day_carry_the_one_core_carries(): + sim = built() + sim.set_input(DRAW, 2024, VALUES) + sim.set_input(DRAW, "year:2024:2", OTHER) + for p in (2025, 2026, 2025): + source = source_of(sim, p) + assert source.period == period("year:2024:2") + np.testing.assert_array_equal(sim.calculate(DRAW, p), OTHER) + np.testing.assert_array_equal(core_answer(sim, p), OTHER) + + +def test_a_register_only_input_uses_its_tiers_mark(): + sim = built({2025: VALUES}) + holder = sim.get_holder(DRAW) + holder._user_input_storage.clear() + assert source_of(sim, 2026).period == period(2025) + np.testing.assert_array_equal(sim.calculate(DRAW, 2026), VALUES) + + +def test_a_dumped_and_restored_draw_is_its_own_source(tmp_path): + from policyengine_core.tools.simulation_dumper import ( + dump_simulation, + restore_simulation, + ) + + sim = built({2025: VALUES}) + dump_simulation(sim, str(tmp_path / "dump")) + restored = restore_simulation(str(tmp_path / "dump"), sim.tax_benefit_system) + source = source_of(restored, 2026) + assert source.period == period(2025) + + +def test_an_after_load_modifier_input_is_supplied_and_its_baseline_lacks_it(): + from policyengine_uk.model_api import Scenario + + def supply(simulation): + simulation.set_input(DRAW, 2025, VALUES) + + sim = Simulation( + dataset=one_person_units_frame(count=3), + scenario=Scenario(simulation_modifier=supply), + ) + assert source_of(sim, 2026).period == period(2025) + np.testing.assert_array_equal(sim.calculate(DRAW, 2026), VALUES) + assert source_of(sim.baseline, 2026) is None + + +def test_a_reform_formula_reads_the_draw_through_the_check(): + sim = built({2025: VALUES}) + + class doubles_the_draw(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "twice the draw" + + def formula(person, period, parameters): + return person(DRAW, period) * 2 + + sim.tax_benefit_system.add_variable(doubles_the_draw) + np.testing.assert_array_equal(sim.calculate("doubles_the_draw", 2026), VALUES * 2) + sim.get_holder(DRAW)._memory_storage.put( + OTHER, period(2027), "default", derived=True + ) + with pytest.raises(StoredDrawError, match="differs from its supplied source"): + sim.calculate("doubles_the_draw", 2027) + missing = built() + missing.tax_benefit_system.add_variable(doubles_the_draw) + np.testing.assert_array_equal( + missing.calculate("doubles_the_draw", 2026), [2, 2, 2] + ) diff --git a/policyengine_uk/tests/test_supplied_input_tightening.py b/policyengine_uk/tests/test_supplied_input_tightening.py index 04ef75178d..eaab3a244d 100644 --- a/policyengine_uk/tests/test_supplied_input_tightening.py +++ b/policyengine_uk/tests/test_supplied_input_tightening.py @@ -37,6 +37,8 @@ supplied_input_periods, ) +pytestmark = pytest.mark.usefixtures("cloned_uk_tax_benefit_system") + NEW_COLUMNS = ( "uc_new_claimant_health_element_draw", "pip_phase_in_draw", diff --git a/policyengine_uk/utils/stored_draws.py b/policyengine_uk/utils/stored_draws.py index e4dc5fadbf..d5fab93370 100644 --- a/policyengine_uk/utils/stored_draws.py +++ b/policyengine_uk/utils/stored_draws.py @@ -256,6 +256,8 @@ def input_source(population, name: str, period) -> Optional[InputSource]: if schema is None: raise StoredDrawError(f"{name} is not a registered stored draw.") simulation = population.simulation + if name not in simulation.tax_benefit_system.variables: + raise StoredDrawError(f"{name} has no current definition.") holder = population.get_holder(name) period = periods.period(period) _check_request(holder, name, schema, period) From d3cd30df8cae76cac9861591a2342b6a84babb6e Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 5 Nov 2026 07:52:04 +0000 Subject: [PATCH 07/11] Test an unregistered name and a source replaced mid-read; drop an unreachable branch Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LwQq6HUE9hj2BPEjEXLL68 --- .../tests/test_stored_draw_input_source.py | 22 +++++++++++++++++++ policyengine_uk/utils/stored_draws.py | 2 -- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/policyengine_uk/tests/test_stored_draw_input_source.py b/policyengine_uk/tests/test_stored_draw_input_source.py index 86e906c52e..9db86203a5 100644 --- a/policyengine_uk/tests/test_stored_draw_input_source.py +++ b/policyengine_uk/tests/test_stored_draw_input_source.py @@ -695,3 +695,25 @@ def formula(person, period, parameters): np.testing.assert_array_equal( missing.calculate("doubles_the_draw", 2026), [2, 2, 2] ) + + +def test_an_unregistered_variable_has_no_draw_source(): + sim = built({2025: VALUES}) + with pytest.raises(StoredDrawError, match="not a registered stored draw"): + input_source(sim.get_variable_population("age"), "age", 2025) + + +def test_a_source_replaced_while_core_reads_it_raises(monkeypatch): + sim = built({2025: VALUES}) + real = CoreSimulation.calculate + + def replaces_the_source(self, name, *args, **kwargs): + result = real(self, name, *args, **kwargs) + if name == DRAW: + # The same values, stored again: a new sequence number. + self.get_holder(DRAW).set_input(period(2025), VALUES) + return result + + monkeypatch.setattr(CoreSimulation, "calculate", replaces_the_source) + with pytest.raises(StoredDrawError, match="changed while"): + sim.calculate(DRAW, 2025) diff --git a/policyengine_uk/utils/stored_draws.py b/policyengine_uk/utils/stored_draws.py index d5fab93370..534a54a92d 100644 --- a/policyengine_uk/utils/stored_draws.py +++ b/policyengine_uk/utils/stored_draws.py @@ -109,8 +109,6 @@ def _describe(name: str, period, branch_name: Optional[str], tier=None) -> str: def _check_schema(variable, name: str, schema: DrawSchema, which: str) -> None: - if variable is None: - raise StoredDrawError(f"{name} has no {which}.") kept = ( variable.entity.key == schema.entity and variable.value_type is schema.value_type From 5a784a7e926d4bc5b1471770ae67d4f96f139140 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 5 Nov 2026 08:01:11 +0000 Subject: [PATCH 08/11] Property test: encoded draws keep every covered decision in both consumers Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LwQq6HUE9hj2BPEjEXLL68 --- .../tests/test_stored_draw_consumers.py | 28 +++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/policyengine_uk/tests/test_stored_draw_consumers.py b/policyengine_uk/tests/test_stored_draw_consumers.py index 817dac9a98..19f58e3f65 100644 --- a/policyengine_uk/tests/test_stored_draw_consumers.py +++ b/policyengine_uk/tests/test_stored_draw_consumers.py @@ -470,3 +470,31 @@ def test_no_take_up_call_passes_a_rate(): if isinstance(node, ast.Name) and node.id == "_GRADUATE_LOAN_TAKE_UP" ] assert len(readers) == 2 # the assignment and the fallback's comparison + + +@settings(PROPERTY) +@given( + legacy=st.lists(st.floats(0, 1, exclude_max=True), min_size=1, max_size=30), + share_index=st.integers(0, len(SHARE_GRID) - 1), +) +def test_encoded_draws_keep_every_grid_decision_in_each_consumer(legacy, share_index): + """Representation: a reference-encoded draw lies in [0, 1), survives a + float32 round trip, sits within 2^-24 of its legacy draw, and gives the + legacy decision in both consumers at every covered threshold.""" + legacy = np.asarray(legacy, dtype=np.float64) + for thresholds, consumer in ( + (SHARE_GRID, new_claimants_by_draw), + (PIP_THRESHOLDS, phased_out_by_draw), + ): + stored, feasible, _ = encode(legacy, thresholds) + assert feasible.all() + as_loaded = np.asarray(stored, dtype=np.float32) + assert ((as_loaded >= 0) & (as_loaded < 1)).all() + np.testing.assert_array_equal( + as_loaded.astype(np.float64).astype(np.float32), as_loaded + ) + assert (np.abs(as_loaded.astype(np.float64) - legacy) < 2.0**-24).all() + threshold = thresholds[share_index % len(thresholds)] + np.testing.assert_array_equal( + consumer(as_loaded, threshold), legacy < threshold + ) From 793468eb3375b2209e586139e4f0eb73e3f1bbbc Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 8 Nov 2026 12:04:34 -0500 Subject: [PATCH 09/11] Protected storage keeps stored draws immutable Co-Authored-By: Claude Opus 5.5 --- .../tests/test_protected_storage.py | 526 ++++++++++++++++++ policyengine_uk/utils/protected_storage.py | 174 ++++++ 2 files changed, 700 insertions(+) create mode 100644 policyengine_uk/tests/test_protected_storage.py create mode 100644 policyengine_uk/utils/protected_storage.py diff --git a/policyengine_uk/tests/test_protected_storage.py b/policyengine_uk/tests/test_protected_storage.py new file mode 100644 index 0000000000..71eaeecdc0 --- /dev/null +++ b/policyengine_uk/tests/test_protected_storage.py @@ -0,0 +1,526 @@ +"""Protected storage keeps values immutable and rejects input changes atomically.""" + +import pickle + +import numpy as np +import pytest +from hypothesis import given, settings +from hypothesis import strategies as st +from policyengine_core.data_storage import InMemoryStorage +from policyengine_core.periods import period + +from policyengine_uk.utils.protected_storage import ProtectedMemoryStorage +from policyengine_uk.utils.stored_draws import StoredDrawError, calculation + + +NAME = "pip_phase_in_draw" +VALUES = np.array([0.25, 0.5, 0.75], dtype=np.float32) +OTHER = np.array([0.125, 0.375, 0.625], dtype=np.float32) + + +def storage(): + result = ProtectedMemoryStorage(False, NAME) + result.put(VALUES, 2025, sequence_number=31) + return result + + +def snapshot(memory): + return ( + { + key: (value.dtype, value.shape, value.tobytes()) + for key, value in memory._arrays.items() + }, + memory._inputs, + memory._sequence_numbers.copy(), + ) + + +def assert_immutable(value, expected): + assert type(value) is np.ndarray + assert isinstance(value.base, bytes) + assert not value.flags.writeable + assert value.dtype == expected.dtype + assert value.shape == expected.shape + assert value.tobytes() == expected.tobytes() + with pytest.raises(ValueError): + value.flags.writeable = True + if value.size: + with pytest.raises(ValueError): + value.flat[0] = 0 + + +@pytest.mark.parametrize("branch", ["default", "child", "ancestor"]) +@pytest.mark.parametrize("target", [2025, 2026, 2027]) +@pytest.mark.parametrize("values", [VALUES, OTHER]) +def test_input_store_refuses_during_a_calculation_without_changing_storage( + branch, target, values +): + memory = storage() + before = snapshot(memory) + with calculation(): + with pytest.raises(StoredDrawError, match=f"{NAME} for {target}.*{branch}"): + memory.put(values, target, branch) + assert snapshot(memory) == before + + +@pytest.mark.parametrize("during", [False, True]) +@pytest.mark.parametrize("values", [VALUES, OTHER]) +def test_derived_store_never_replaces_a_stored_input(during, values): + memory = storage() + before = snapshot(memory) + + def attempt(): + with pytest.raises(StoredDrawError, match="over a stored input"): + memory.put(values, 2025, derived=True) + + if during: + with calculation(): + attempt() + else: + attempt() + assert snapshot(memory) == before + + +def test_derived_values_can_change_during_a_calculation(): + memory = storage() + with calculation(): + memory.put(VALUES, 2026, derived=True, sequence_number=32) + memory.put(OTHER, 2026, derived=True, sequence_number=33) + assert memory.is_derived(2026) + assert_immutable(memory.get(2026), OTHER) + assert memory._sequence_numbers["default:2026"] == 33 + assert memory._inputs == frozenset({"default:2025"}) + + +def test_an_input_can_replace_a_derived_value_between_calculations(): + memory = storage() + memory.put(VALUES, 2026, derived=True) + memory.put(OTHER, 2026) + assert not memory.is_derived(2026) + assert_immutable(memory.get(2026), OTHER) + assert isinstance(memory._inputs, frozenset) + + +@pytest.mark.parametrize("target", [None, 2025, "year:2024:3"]) +def test_delete_refuses_before_removing_any_selected_value(target): + memory = storage() + memory.put(OTHER, 2024, derived=True) + memory.put(OTHER, 2026, derived=True) + before = snapshot(memory) + with calculation(): + with pytest.raises(StoredDrawError, match=f"{NAME} for 2025.*default"): + memory.delete(target) + assert snapshot(memory) == before + + +def test_delete_and_cache_drop_keep_inputs_while_removing_only_derived_values(): + memory = storage() + memory.put(OTHER, 2025, "child", derived=True) + memory.put(OTHER, 2026, derived=True) + memory.put(OTHER, 2027, derived=True) + with calculation(): + memory.delete(2025, "child") + memory.delete(2026) + memory.mark_derived_except({"default:2025"}) + assert memory.drop_computed() == 1 + assert_immutable(memory.get(2025), VALUES) + assert memory._inputs == frozenset({"default:2025"}) + assert memory._sequence_numbers == {"default:2025": 31} + + +@pytest.mark.parametrize("operation", ["mark", "unmark"]) +def test_actual_input_demotion_refuses_without_changing_any_mark(operation): + memory = storage() + memory.put(OTHER, 2026) + before = snapshot(memory) + with calculation(): + with pytest.raises(StoredDrawError, match=f"demote an input of {NAME}"): + if operation == "mark": + memory.mark_derived_except(set()) + else: + memory._unmark_input("default:2025") + assert snapshot(memory) == before + + +def test_unchanged_marks_and_absent_marks_pass_during_a_calculation(): + memory = storage() + before = snapshot(memory) + with calculation(): + memory.mark_derived_except({"default:2025", "default:2026"}) + memory._unmark_input("default:2026") + assert snapshot(memory) == before + + +def test_a_one_shot_kept_key_iterable_cannot_demote_an_input(): + memory = storage() + before = snapshot(memory) + with calculation(): + memory.mark_derived_except(iter(["default:2025"])) + assert snapshot(memory) == before + + +@pytest.mark.parametrize("during", [False, True]) +def test_a_derived_store_clears_only_its_stale_input_mark(during): + memory = storage() + memory._inputs = memory._inputs | {"default:2026", "default:2027"} + + def store_derived(): + memory.put(OTHER, 2026, derived=True, sequence_number=32) + + if during: + with calculation(): + store_derived() + else: + store_derived() + assert memory._inputs == frozenset({"default:2025", "default:2027"}) + assert memory.is_derived(2026) + assert_immutable(memory.get(2025), VALUES) + assert_immutable(memory.get(2026), OTHER) + assert memory._sequence_numbers == {"default:2025": 31, "default:2026": 32} + + +@pytest.mark.parametrize("during", [False, True]) +@pytest.mark.parametrize("invalid", ["array", "branch", "period"]) +def test_an_invalid_derived_store_keeps_stale_marks_and_storage_unchanged( + during, invalid +): + memory = storage() + values, target, branch = OTHER, 2026, "default" + if invalid == "array": + values = OTHER.astype(object) + elif invalid == "branch": + branch = "bad:branch" + else: + target = "year:2026-01-02" + memory._inputs = memory._inputs | {memory._key(target, branch)} + before = snapshot(memory) + + def attempt(): + with pytest.raises(ValueError): + memory.put(values, target, branch, derived=True) + + if during: + with calculation(): + attempt() + else: + attempt() + assert snapshot(memory) == before + + +def test_between_calculation_deletions_and_demotions_keep_frozen_marks(): + memory = storage() + memory.put(OTHER, 2026) + memory.mark_derived_except({"default:2025"}) + assert isinstance(memory._inputs, frozenset) + memory.delete(2026) + assert isinstance(memory._inputs, frozenset) + memory.delete() + assert memory._inputs == frozenset() + assert memory._sequence_numbers == {} + memory.put(VALUES, 2025) + memory._unmark_input("default:2025") + assert memory._inputs == frozenset() + + +@pytest.mark.parametrize("derived", [False, True]) +def test_put_copies_the_callers_memory_and_get_returns_a_separate_array(derived): + memory = ProtectedMemoryStorage(False, NAME) + given = VALUES.copy() + memory.put(given, 2025, derived=derived) + given[:] = 0 + first = memory.get(2025) + second = memory.get(2025) + assert first is not second + assert first is not memory._arrays["default:2025"] + assert first.base is second.base + assert_immutable(first, VALUES) + first.shape = (3, 1) + second.dtype = np.uint32 + assert_immutable(memory.get(2025), VALUES) + + +@pytest.mark.parametrize( + "values", + [np.array(0.25), np.arange(6).reshape(2, 3), np.empty((0, 3)), np.arange(10)[::2]], +) +def test_every_shape_and_layout_reads_over_bytes(values): + memory = ProtectedMemoryStorage(False, NAME) + memory.put(values, 2025) + assert_immutable(memory.get(2025), values) + assert isinstance(memory._arrays["default:2025"].base, bytes) + + +@pytest.mark.parametrize( + "values", + [ + np.array([0.0, -0.0, np.inf, -np.inf, np.nan], dtype=">f8"), + np.array([complex(-0.0, 1), complex(2, -0.0), complex(np.nan, np.inf)]), + ], +) +def test_numeric_special_values_keep_their_stored_bits(values): + memory = ProtectedMemoryStorage(False, NAME) + memory.put(values, 2025) + assert_immutable(memory.get(2025), values) + assert_immutable(memory.clone().get(2025), values) + assert_immutable(pickle.loads(pickle.dumps(memory)).get(2025), values) + + +class ArraySubclass(np.ndarray): + pass + + +@pytest.mark.parametrize( + "values", + [ + np.ma.array(VALUES, mask=[False, True, False]), + VALUES.view(ArraySubclass), + np.array([1, 2, 3], dtype=object), + np.array(["1", "2", "3"]), + np.array([b"1", b"2", b"3"]), + np.array(["2025-01-01"], dtype="datetime64[D]"), + np.zeros(3, dtype=[("draw", "f4")]), + [0.25, 0.5, 0.75], + ], +) +def test_non_plain_or_non_numeric_values_refuse_before_changing_storage(values): + memory = storage() + before = snapshot(memory) + with pytest.raises(StoredDrawError, match=NAME): + memory.put(values, 2026) + assert snapshot(memory) == before + + +@pytest.mark.parametrize("share_arrays", [False, True]) +def test_clone_shares_only_immutable_arrays_and_copies_indices(share_arrays): + memory = storage() + memory.put(OTHER, 2026, derived=True, sequence_number=32) + child = memory.clone(share_arrays=share_arrays) + assert isinstance(child, ProtectedMemoryStorage) + assert snapshot(child) == snapshot(memory) + assert child._arrays is not memory._arrays + assert child._arrays["default:2025"] is memory._arrays["default:2025"] + assert child._sequence_numbers is not memory._sequence_numbers + assert isinstance(child._inputs, frozenset) + returned = child.get(2025) + returned.dtype = np.uint8 + assert_immutable(child.get(2025), VALUES) + assert_immutable(memory.get(2025), VALUES) + child.put(OTHER, 2025) + child.delete(2026) + assert_immutable(memory.get(2025), VALUES) + assert_immutable(memory.get(2026), OTHER) + + +def test_getstate_hands_out_no_mutable_container_or_kept_array(): + memory = storage() + before = snapshot(memory) + state = memory.__getstate__() + assert state["_arrays"] is not memory._arrays + assert state["_arrays"]["default:2025"] is not memory._arrays["default:2025"] + assert state["_sequence_numbers"] is not memory._sequence_numbers + assert "_array_dtypes" not in memory.__dict__ + assert isinstance(state["_inputs"], frozenset) + state["_arrays"]["default:2025"].dtype = np.uint8 + state["_arrays"].clear() + state["_sequence_numbers"].clear() + state["_array_dtypes"].clear() + state["_inputs"] = frozenset() + assert snapshot(memory) == before + + +def test_standard_pickle_restores_every_value_over_immutable_bytes(): + memory = storage() + memory.put(OTHER, 2026, derived=True, sequence_number=32) + with calculation(): + restored = pickle.loads(pickle.dumps(memory)) + assert snapshot(restored) == snapshot(memory) + assert isinstance(restored._inputs, frozenset) + assert_immutable(restored.get(2025), VALUES) + assert_immutable(restored.get(2026), OTHER) + + +def test_pickle_keeps_an_empty_non_native_array_dtype(): + values = np.empty(0, dtype=">i2") + memory = ProtectedMemoryStorage(False, NAME) + memory.put(values, 2025) + restored = pickle.loads(pickle.dumps(memory)) + assert_immutable(restored.get(2025), values) + + +def test_pickling_exposes_only_detached_array_objects(): + import io + + memory = storage() + arrays = [] + + class ArrayCollector(pickle.Pickler): + def persistent_id(self, value): + if isinstance(value, np.ndarray): + arrays.append(value) + return None + + ArrayCollector(io.BytesIO()).dump(memory) + assert arrays + for value in arrays: + assert value is not memory._arrays["default:2025"] + value.shape = (value.size, 1) + value.dtype = np.uint8 + assert_immutable(memory.get(2025), VALUES) + assert_immutable(memory.clone().get(2025), VALUES) + + +def test_setstate_copies_legacy_values_and_marks_without_changing_the_state(): + memory = ProtectedMemoryStorage(False, NAME) + given = VALUES.copy() + state = {"is_eternal": False, "_arrays": {"default:2025": given}} + memory.__setstate__(state) + assert state["_arrays"]["default:2025"] is given + given[:] = 0 + assert_immutable(memory.get(2025), VALUES) + assert memory._inputs == frozenset({"default:2025"}) + assert memory._sequence_numbers == {} + + +def test_from_storage_preserves_values_marks_and_sequence_numbers(): + plain = InMemoryStorage(False) + plain.put(VALUES.copy(), 2025, sequence_number=31) + plain.put(OTHER.copy(), 2026, derived=True, sequence_number=32) + protected = ProtectedMemoryStorage.from_storage(plain, NAME) + assert snapshot(protected) == snapshot(plain) + assert protected.variable_name == NAME + plain.get(2025)[:] = 0 + assert_immutable(protected.get(2025), VALUES) + assert_immutable(protected.get(2026), OTHER) + assert isinstance(protected._inputs, frozenset) + + +def test_from_storage_copies_a_shared_array_without_changing_the_source(): + source = InMemoryStorage(False) + source.put(VALUES.copy(), 2025, sequence_number=31) + plain = source.clone(share_arrays=True) + before = snapshot(plain) + protected = ProtectedMemoryStorage.from_storage(plain, NAME) + assert snapshot(plain) == before + source.get(2025)[:] = 0 + assert_immutable(protected.get(2025), VALUES) + + +def test_legacy_state_keeps_its_recorded_derived_values_derived(): + memory = ProtectedMemoryStorage(False, NAME) + memory.__setstate__( + { + "is_eternal": False, + "_arrays": {"default:2025": VALUES.copy(), "default:2026": OTHER.copy()}, + "_derived": {"default:2026"}, + } + ) + assert memory._inputs == frozenset({"default:2025"}) + assert memory.is_derived(2026) + assert_immutable(memory.get(2025), VALUES) + assert_immutable(memory.get(2026), OTHER) + + +def test_eternal_keys_and_refusals_follow_the_storage_period(): + memory = ProtectedMemoryStorage(True, NAME) + memory.put(VALUES, 2025) + assert_immutable(memory.get(2040), VALUES) + assert memory._inputs == frozenset({"default:ETERNITY"}) + with calculation(): + with pytest.raises(StoredDrawError): + memory.delete(2040) + with pytest.raises(StoredDrawError): + memory.put(OTHER, 2040, derived=True) + assert_immutable(memory.get(2025), VALUES) + + +def test_core_key_validation_still_precedes_any_storage_change(): + memory = storage() + before = snapshot(memory) + with pytest.raises(ValueError, match="may not contain"): + memory.put(VALUES, 2026, "bad:branch") + with pytest.raises(ValueError, match="mid-month"): + memory.put(VALUES, "year:2026-01-02") + assert snapshot(memory) == before + + +NUMERIC_DTYPES = tuple( + dict.fromkeys( + np.dtype(dtype).newbyteorder(order) + for dtype in [ + np.bool_, + np.int8, + np.int16, + np.int32, + np.int64, + np.uint8, + np.uint16, + np.uint32, + np.uint64, + np.float16, + np.float32, + np.float64, + np.longdouble, + np.complex64, + np.complex128, + np.clongdouble, + ] + for order in ["<", ">"] + ) +) + + +@settings(max_examples=150, deadline=None) +@given( + dtype=st.sampled_from(NUMERIC_DTYPES), + size=st.integers(0, 64), + data=st.data(), +) +def test_numeric_dtypes_and_sizes_round_trip_over_immutable_bytes(dtype, size, data): + byte_count = size * dtype.itemsize + raw = data.draw(st.binary(min_size=byte_count, max_size=byte_count)) + given = np.frombuffer(raw, dtype=dtype) + memory = ProtectedMemoryStorage(False, NAME) + memory.put(given, 2025) + assert_immutable(memory.get(2025), given) + assert_immutable(memory.clone().get(2025), given) + assert_immutable(pickle.loads(pickle.dumps(memory)).get(2025), given) + + +@settings(max_examples=100, deadline=None) +@given( + operations=st.lists( + st.sampled_from( + ["read", "clone", "derived", "delete", "drop", "marks", "input", "demote"] + ), + max_size=40, + ) +) +def test_no_sequence_of_operations_during_a_calculation_changes_an_input(operations): + memory = storage() + input_record = (memory._inputs, memory._sequence_numbers["default:2025"]) + with calculation(): + for operation in operations: + if operation == "read": + returned = memory.get(2025) + returned.shape = (1, 3) + elif operation == "clone": + assert_immutable(memory.clone().get(2025), VALUES) + elif operation == "derived": + memory.put(OTHER, 2026, derived=True) + elif operation == "delete": + memory.delete(2026) + elif operation == "drop": + memory.drop_computed() + elif operation == "marks": + memory.mark_derived_except({"default:2025"}) + elif operation == "input": + with pytest.raises(StoredDrawError): + memory.put(OTHER, 2025) + else: + with pytest.raises(StoredDrawError): + memory.mark_derived_except(set()) + assert_immutable(memory.get(2025), VALUES) + assert ( + memory._inputs, + memory._sequence_numbers["default:2025"], + ) == input_record diff --git a/policyengine_uk/utils/protected_storage.py b/policyengine_uk/utils/protected_storage.py new file mode 100644 index 0000000000..3a4ea0d5a2 --- /dev/null +++ b/policyengine_uk/utils/protected_storage.py @@ -0,0 +1,174 @@ +"""Memory storage that keeps stored draws over immutable bytes.""" + +import numpy as np +from policyengine_core import periods +from policyengine_core.data_storage import InMemoryStorage + +from policyengine_uk.utils.stored_draws import ( + StoredDrawError, + calculation_in_flight, +) + + +class ProtectedMemoryStorage(InMemoryStorage): + """Keep numeric values immutable and refuse input changes during a calculation.""" + + def __init__(self, is_eternal: bool, variable_name: str): + super().__init__(is_eternal) + self.variable_name = variable_name + self._inputs = frozenset() + + def _immutable_array(self, value): + if type(value) is not np.ndarray: + raise StoredDrawError( + f"{self.variable_name} needs a plain numeric numpy.ndarray; " + f"{type(value).__name__} cannot supply protected storage." + ) + if value.dtype.kind not in "biufc": + raise StoredDrawError( + f"{self.variable_name} needs a numeric dtype; " + f"{value.dtype} cannot supply protected storage." + ) + result = np.frombuffer(value.tobytes(), dtype=value.dtype) + if result.shape != value.shape: + result.shape = value.shape + return result + + @staticmethod + def _new_array(value): + result = np.frombuffer(value.base, dtype=value.dtype) + if result.shape != value.shape: + result.shape = value.shape + return result + + def _key(self, period, branch_name): + if self.is_eternal: + period = periods.period(periods.ETERNITY) + return f"{branch_name}:{periods.period(period)}" + + def _refuse(self, action, period, branch_name): + raise StoredDrawError( + f"Cannot {action} {self.variable_name} for {period} " + f"on branch {branch_name!r} during a calculation." + ) + + def put( + self, + value, + period, + branch_name="default", + derived=False, + sequence_number=None, + ): + if not derived and calculation_in_flight(): + self._refuse("store an input of", period, branch_name) + key = self._key(period, branch_name) + if derived and key in self._arrays and key in self._inputs: + raise StoredDrawError( + f"Cannot store a derived value of {self.variable_name} for " + f"{period} on branch {branch_name!r} over a stored input." + ) + value = self._immutable_array(value) + inputs = self._inputs + if derived and key not in self._arrays: + self._inputs = inputs - {key} + try: + super().put( + value, + period, + branch_name, + derived=derived, + sequence_number=sequence_number, + ) + except Exception: + self._inputs = inputs + raise + + def get(self, period, branch_name="default"): + value = self._arrays.get(self._key(period, branch_name)) + return None if value is None else self._new_array(value) + + def delete(self, period=None, branch_name="default"): + prefix = f"{branch_name}:" + if period is not None: + period = periods.period(periods.ETERNITY if self.is_eternal else period) + selected = { + key + for key in self._arrays + if key.startswith(prefix) + and ( + period is None or period.contains(periods.period(key.split(":", 1)[1])) + ) + } + selected_inputs = selected.intersection(self._inputs) + if calculation_in_flight() and selected_inputs: + stored_branch, stored_period = next(iter(selected_inputs)).split(":", 1) + self._refuse("delete an input of", stored_period, stored_branch) + super().delete(period, branch_name) + + def _mark_input(self, key): + self._inputs = self._inputs | {key} + + def _unmark_input(self, key): + if calculation_in_flight() and key in self._inputs and key in self._arrays: + branch_name, period = key.split(":", 1) + self._refuse("demote an input of", period, branch_name) + self._inputs = self._inputs - {key} + + def _unmark_dropped_keys(self): + self._inputs = self._inputs.intersection(self._arrays) + + def mark_derived_except(self, keys): + keys = frozenset(keys) + demoted = self._inputs.difference(keys) + if calculation_in_flight() and demoted.intersection(self._arrays): + key = next(iter(demoted.intersection(self._arrays))) + branch_name, period = key.split(":", 1) + self._refuse("demote an input of", period, branch_name) + self._inputs = self._inputs.intersection(keys) + + def clone(self, share_arrays=False): + """Copy the indices and share only immutable arrays and input marks.""" + result = type(self)(self.is_eternal, self.variable_name) + result._arrays = self._arrays.copy() + result._inputs = self._inputs.intersection(result._arrays) + result._sequence_numbers = self._sequence_numbers.copy() + return result + + @classmethod + def from_storage(cls, storage, variable_name): + """Copy a storage's values, input marks and sequence numbers.""" + result = cls(storage.is_eternal, variable_name) + state = storage.__getstate__() + state["variable_name"] = variable_name + result.__setstate__(state) + return result + + def __getstate__(self): + state = super().__getstate__() + state["_arrays"] = { + key: self._new_array(value) for key, value in self._arrays.items() + } + state["_inputs"] = frozenset(self._inputs) + state["_sequence_numbers"] = self._sequence_numbers.copy() + state["_array_dtypes"] = { + key: value.dtype.str for key, value in self._arrays.items() + } + state.pop("_shared", None) + return state + + def __setstate__(self, state): + state = state.copy() + dtypes = state.pop("_array_dtypes", {}) + super().__setstate__(state) + arrays = {} + for key, value in self._arrays.items(): + dtype = np.dtype(dtypes.get(key, value.dtype)) + if value.dtype != dtype: + # NumPy's standard pickler restores numeric arrays in native byte order. + value = value.byteswap().view(dtype) + arrays[key] = self._immutable_array(value) + self._arrays = arrays + self._inputs = frozenset(self._inputs) + self._sequence_numbers = self._sequence_numbers.copy() + self.__dict__.pop("_shared", None) From c1e4da397a34dafd304b8d4201992d6da11e285e Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 8 Nov 2026 12:04:34 -0500 Subject: [PATCH 10/11] Draw holders protect inputs before data loads Co-Authored-By: Claude Opus 5.5 --- policyengine_uk/simulation.py | 18 +- .../tests/test_protected_draw_storage.py | 452 ++++++++++++++++++ .../tests/test_protected_storage.py | 29 ++ .../tests/test_stored_draw_input_source.py | 29 +- policyengine_uk/utils/protected_inputs.py | 86 ++++ policyengine_uk/utils/protected_storage.py | 4 +- policyengine_uk/utils/stored_draws.py | 7 + 7 files changed, 617 insertions(+), 8 deletions(-) create mode 100644 policyengine_uk/tests/test_protected_draw_storage.py create mode 100644 policyengine_uk/utils/protected_inputs.py diff --git a/policyengine_uk/simulation.py b/policyengine_uk/simulation.py index c6b71a3f1e..7ff11f71e3 100644 --- a/policyengine_uk/simulation.py +++ b/policyengine_uk/simulation.py @@ -113,6 +113,7 @@ check_option_request, input_source, ) +from policyengine_uk.utils.protected_inputs import install_draw_storage, protects_draws from policyengine_uk.reforms import create_structural_reforms_from_parameters from policyengine_uk.parameters.gov.simulation.labour_supply_responses.aliases import ( canonicalize_parameter_changes, @@ -651,6 +652,7 @@ def build_from_populations(self, populations) -> None: # the values the new populations hold, and only those: an equal value # a rebuild stores later (from a situation or a frame) is supplied. keep_population_records(self, adopted) + install_draw_storage(self) def set_input(self, variable_name: str, period: Period, value) -> None: if variable_name in REGISTERED_DRAW_NAMES: @@ -1833,7 +1835,9 @@ def calculate( drop_missing_supplied_inputs(self, variable_name) draw_source = _UNCHECKED - if variable_name in REGISTERED_DRAW_NAMES and built_from_data(self): + if variable_name in REGISTERED_DRAW_NAMES and ( + protects_draws(self) or built_from_data(self) + ): # A stored draw: decide which supplied input core reads, after the # cleanups above and before core, and never answer from the fast # cache, which can outlive the value the check examined. @@ -1889,7 +1893,9 @@ def _calculate_source_backed_draw(self, variable_name, period, source, map_to): return raw def calculate_add(self, variable_name: str, period: Period = None, *args, **kwargs): - if variable_name in REGISTERED_DRAW_NAMES and built_from_data(self): + if variable_name in REGISTERED_DRAW_NAMES and ( + protects_draws(self) or built_from_data(self) + ): check_option_request(self, variable_name, "calculate_add") with calculation(): return super().calculate_add(variable_name, period, *args, **kwargs) @@ -1897,13 +1903,17 @@ def calculate_add(self, variable_name: str, period: Period = None, *args, **kwar def calculate_divide( self, variable_name: str, period: Period = None, *args, **kwargs ): - if variable_name in REGISTERED_DRAW_NAMES and built_from_data(self): + if variable_name in REGISTERED_DRAW_NAMES and ( + protects_draws(self) or built_from_data(self) + ): check_option_request(self, variable_name, "calculate_divide") with calculation(): return super().calculate_divide(variable_name, period, *args, **kwargs) def calculate_output(self, variable_name: str, period: Period = None): - if variable_name in REGISTERED_DRAW_NAMES and built_from_data(self): + if variable_name in REGISTERED_DRAW_NAMES and ( + protects_draws(self) or built_from_data(self) + ): check_option_request(self, variable_name, "calculate_output") with calculation(): return super().calculate_output(variable_name, period) diff --git a/policyengine_uk/tests/test_protected_draw_storage.py b/policyengine_uk/tests/test_protected_draw_storage.py new file mode 100644 index 0000000000..c8377b6e35 --- /dev/null +++ b/policyengine_uk/tests/test_protected_draw_storage.py @@ -0,0 +1,452 @@ +"""Data attachments retain immutable draw storage and verify every visible input.""" + +import copy + +import numpy as np +import pytest +from policyengine_core.data_storage import InMemoryStorage, OnDiskStorage +from policyengine_core.holders import Holder +from policyengine_core.periods import period + +from policyengine_uk import Simulation +from policyengine_uk.model_api import Person, Variable, YEAR +from policyengine_uk.tests.stored_draw_fixtures import one_person_units_frame +from policyengine_uk.tests.test_stored_draw_core_route import core_file +from policyengine_uk.utils.protected_inputs import install_draw_storage, protects_draws +from policyengine_uk.utils.protected_storage import ProtectedMemoryStorage +from policyengine_uk.utils.stored_draws import ( + REGISTERED_DRAW_NAMES, + StoredDrawError, + calculation, + input_source, +) + + +pytestmark = pytest.mark.usefixtures("cloned_uk_tax_benefit_system") + +DRAW = "pip_phase_in_draw" +GROUP_DRAW = "uc_new_claimant_health_element_draw" +VALUES = np.array([0.25, 0.5, 0.75], dtype=np.float32) +OTHER = np.array([0.125, 0.375, 0.625], dtype=np.float32) + + +def built(columns=None): + if columns is None: + columns = {DRAW: VALUES, GROUP_DRAW: OTHER} + return Simulation(dataset=one_person_units_frame(columns=columns)) + + +def situation(): + names = [f"person_{index}" for index in range(len(VALUES))] + return { + "people": { + name: {"age": 30, DRAW: float(values)} + for name, values in zip(names, VALUES) + }, + "benunits": { + name: {"members": [name], GROUP_DRAW: float(values)} + for name, values in zip(names, OTHER) + }, + "households": {name: {"members": [name]} for name in names}, + } + + +def source_of(simulation, requested=2026, name=DRAW): + return input_source(simulation.get_variable_population(name), name, requested) + + +def snapshot(simulation, name=DRAW): + holder = simulation.get_holder(name) + storage = holder._memory_storage + return ( + { + key: (values.dtype, values.shape, values.tobytes()) + for key, values in storage._arrays.items() + }, + frozenset(storage._inputs), + storage._sequence_numbers.copy(), + set(simulation._user_input_keys), + copy.deepcopy(getattr(holder, "_user_input_storage", {})), + ) + + +def assert_immutable(actual, expected): + assert type(actual) is np.ndarray + assert isinstance(actual.base, bytes) + assert not actual.flags.writeable + np.testing.assert_array_equal(actual, expected) + with pytest.raises(ValueError): + actual.flags.writeable = True + with pytest.raises(ValueError): + actual[0] = 0 + + +@pytest.mark.parametrize("route", ["frame", "core"]) +def test_each_loader_installs_both_draw_storages_before_its_first_store( + route, monkeypatch, tmp_path +): + seen = [] + original = Holder._set + + def checked_store(holder, *args, **kwargs): + if holder.variable.name in REGISTERED_DRAW_NAMES: + seen.append(holder.variable.name) + assert isinstance(holder._memory_storage, ProtectedMemoryStorage) + return original(holder, *args, **kwargs) + + monkeypatch.setattr(Holder, "_set", checked_store) + if route == "frame": + simulation = built() + else: + path = core_file( + tmp_path / "draws.h5", + columns={DRAW: {2025: VALUES}, GROUP_DRAW: {2025: OTHER}}, + ) + simulation = Simulation(dataset=path) + + assert set(seen) == REGISTERED_DRAW_NAMES + assert protects_draws(simulation) + for name, expected in ((DRAW, VALUES), (GROUP_DRAW, OTHER)): + assert_immutable(simulation.calculate(name, 2025), expected) + + +def test_data_attachment_installs_empty_registered_holders(): + simulation = built({}) + for name in REGISTERED_DRAW_NAMES: + assert isinstance( + simulation.get_holder(name)._memory_storage, ProtectedMemoryStorage + ) + + +def test_the_installer_does_not_create_holders_on_a_situation(): + from policyengine_uk.system import system + + simulation = Simulation.__new__(Simulation) + simulation.tax_benefit_system = system + simulation.populations = system.instantiate_entities() + simulation.built_from_dataset = False + install_draw_storage(simulation) + assert not any( + name in population._holders + for population in simulation.populations.values() + for name in REGISTERED_DRAW_NAMES + ) + + +def test_root_clone_and_branch_keep_input_marks_values_and_sequence_numbers(): + simulation = built() + assert_immutable(simulation.calculate(DRAW, 2026), VALUES) + before = snapshot(simulation) + source = source_of(simulation) + + clone = simulation.clone() + branch = simulation.get_branch("draw_child") + for target in (simulation, clone, branch): + assert protects_draws(target) + assert snapshot(target) == before + assert source_of(target) == source + assert_immutable(target.calculate(DRAW, 2025), VALUES) + first = target.calculate(DRAW, 2026) + second = target.calculate(DRAW, 2026) + assert first is not second + assert_immutable(first, VALUES) + first.shape = (1, 3) + first.dtype = np.uint8 + assert_immutable(target.calculate(DRAW, 2026), VALUES) + + assert clone._user_input_keys is not simulation._user_input_keys + assert branch._user_input_keys is not simulation._user_input_keys + root_storage = simulation.get_holder(DRAW)._memory_storage + for target in (clone, branch): + storage = target.get_holder(DRAW)._memory_storage + assert storage is not root_storage + assert storage._arrays is not root_storage._arrays + assert storage._sequence_numbers is not root_storage._sequence_numbers + + +def test_a_data_rebuild_adopts_plain_inputs_without_changing_their_provenance(): + source = Simulation(situation=situation()) + source_holder = source.get_holder(DRAW) + before = snapshot(source) + simulation = built() + simulation.build_from_populations(source.populations) + + assert snapshot(simulation) == before + assert source_of(simulation, 2025).sequence_number == before[2]["default:2025"] + assert_immutable(simulation.calculate(DRAW, 2025), VALUES) + source_holder.get_array(period(2025))[0] = 0 + assert_immutable(simulation.calculate(DRAW, 2025), VALUES) + assert simulation.populations["person"] is not source.populations["person"] + assert simulation.get_holder(DRAW) is not source_holder + + +def test_a_same_population_rebuild_keeps_its_existing_draw_record(): + simulation = built() + before = snapshot(simulation) + simulation.build_from_populations(simulation.populations) + assert snapshot(simulation) == before + assert_immutable(simulation.calculate(DRAW, 2026), VALUES) + + +def test_a_rebuild_with_fresh_populations_installs_before_new_inputs(): + simulation = built() + populations = simulation.tax_benefit_system.instantiate_entities() + for population in populations.values(): + population.count = 3 + population.ids = np.arange(3) + if not population.entity.is_person: + population.members_entity_id = np.arange(3) + population.members_role = np.full(3, population.entity.flattened_roles[0]) + simulation.build_from_populations(populations) + + for name in REGISTERED_DRAW_NAMES: + holder = simulation.get_holder(name) + assert isinstance(holder._memory_storage, ProtectedMemoryStorage) + assert holder.get_known_periods() == [] + assert not any(key[0] == name for key in simulation._user_input_keys) + simulation.set_input(DRAW, 2025, VALUES) + assert_immutable(simulation.calculate(DRAW, 2026), VALUES) + + +@pytest.mark.parametrize("replacement", ["holder", "storage"]) +def test_plain_replacements_fail_every_read_and_a_clone_cannot_trust_them(replacement): + simulation = built() + population = simulation.get_variable_population(DRAW) + holder = population.get_holder(DRAW) + if replacement == "holder": + holder = Holder(holder.variable, population) + population._holders[DRAW] = holder + else: + holder._memory_storage = InMemoryStorage(False) + holder._memory_storage.put(VALUES.copy(), period(2025)) + + for _ in range(2): + with pytest.raises(StoredDrawError, match=f"{DRAW}.*protected memory"): + simulation.calculate(DRAW, 2025) + try: + clone = simulation.clone() + except StoredDrawError: + return + for _ in range(2): + with pytest.raises(StoredDrawError, match=f"{DRAW}.*protected memory"): + clone.calculate(DRAW, 2025) + + +@pytest.mark.parametrize("record", ["missing_key", "none"]) +def test_a_visible_mark_requires_the_reading_simulations_record(record): + simulation = built() + if record == "missing_key": + simulation._user_input_keys.discard((DRAW, "default", period(2025))) + else: + simulation._user_input_keys = None + with pytest.raises(StoredDrawError, match="simulation's input record"): + source_of(simulation, 2025) + + +@pytest.mark.parametrize("requested", [2024, 2027]) +def test_a_later_visible_mark_cannot_disappear_from_the_record(requested): + simulation = built() + simulation.get_holder(DRAW).set_input(period(2024), OTHER) + simulation._user_input_keys.remove((DRAW, "default", period(2025))) + with pytest.raises(StoredDrawError, match=f"{DRAW} for 2025.*input record"): + source_of(simulation, requested) + + +def test_a_shadowed_visible_parent_mark_still_requires_a_record(): + simulation = built() + branch = simulation.get_branch("record_child") + branch.get_holder(DRAW).set_input(period(2025), OTHER, branch.branch_name) + branch._user_input_keys.remove((DRAW, "default", period(2025))) + with pytest.raises( + StoredDrawError, match=f"{DRAW} for 2025.*default.*input record" + ): + source_of(branch, 2025) + assert_immutable(simulation.calculate(DRAW, 2025), VALUES) + + +def test_a_mark_outside_the_readable_branches_does_not_block_a_read(): + simulation = built() + simulation.get_holder(DRAW)._memory_storage.put( + OTHER.copy(), period(2024), "unread_branch" + ) + assert source_of(simulation, 2025).period == period(2025) + assert_immutable(simulation.calculate(DRAW, 2025), VALUES) + + +@pytest.mark.parametrize( + "location", + [ + None, + (period(2024), frozenset({"memory"})), + (period(2025), frozenset()), + (period(2025), frozenset({"disk"})), + ], + ids=["none", "wrong_period", "empty", "wrong_tier"], +) +def test_a_visible_mark_requires_the_holders_matching_location_record(location): + simulation = built() + holder = simulation.get_holder(DRAW) + holder._user_input_storage[("default", "2025")] = location + with pytest.raises(StoredDrawError, match="holder's location record"): + source_of(simulation, 2025) + + +def test_a_register_only_input_seeds_its_location_without_changing_storage(): + simulation = built() + holder = simulation.get_holder(DRAW) + del holder._user_input_storage[("default", "2025")] + before = snapshot(simulation) + + source = source_of(simulation, 2025) + + assert source.period == period(2025) + assert snapshot(simulation)[:4] == before[:4] + assert holder._user_input_storage[("default", "2025")] == ( + period(2025), + frozenset({"memory"}), + ) + + +def test_an_early_year_mark_keeps_its_period_when_the_storage_key_omits_zeroes(): + simulation = built() + early = period("0001") + holder = simulation.get_holder(DRAW) + holder.set_input(early, OTHER) + source = source_of(simulation, early) + assert source.period == early + assert ( + source.sequence_number == holder._memory_storage._sequence_numbers["default:1"] + ) + assert_immutable(holder.get_array(early), OTHER) + + +def test_clearing_the_public_data_flag_keeps_protection_and_clone_protection(): + simulation = built() + simulation.built_from_dataset = False + clone = simulation.clone() + for target in (simulation, clone): + assert protects_draws(target) + assert_immutable(target.calculate(DRAW, 2025), VALUES) + target._user_input_keys.remove((DRAW, "default", period(2025))) + with pytest.raises(StoredDrawError, match="simulation's input record"): + target.calculate(DRAW, 2025) + + +@pytest.mark.parametrize( + "method", ["calculate_add", "calculate_divide", "calculate_output"] +) +def test_clearing_the_public_data_flag_keeps_draw_option_refusals(method): + simulation = built() + simulation.built_from_dataset = False + with pytest.raises(StoredDrawError, match=method): + getattr(simulation, method)(DRAW, 2025) + + +def test_an_explicit_situation_rebuild_replaces_protected_storage_with_plain_storage(): + simulation = built() + simulation.build_from_situation(situation()) + assert not protects_draws(simulation) + holder = simulation.get_holder(DRAW) + assert type(holder._memory_storage) is InMemoryStorage + with calculation(): + holder.set_input(period(2026), OTHER) + np.testing.assert_array_equal(simulation.calculate(DRAW, 2026), OTHER) + + +def test_a_situation_adopts_data_inputs_as_ordinary_writeable_storage(): + source = built() + before = snapshot(source) + simulation = Simulation(situation=situation()) + simulation.build_from_populations(source.populations) + assert not protects_draws(simulation) + holder = simulation.get_holder(DRAW) + assert type(holder._memory_storage) is InMemoryStorage + assert holder.get_array(period(2025)).flags.writeable + with calculation(): + holder.set_input(period(2026), OTHER) + np.testing.assert_array_equal(simulation.calculate(DRAW, 2026), OTHER) + assert snapshot(source) == before + + +@pytest.mark.parametrize("target", ["own", "unrelated"]) +@pytest.mark.parametrize("operation", ["input", "delete", "demote", "unmark"]) +def test_a_formula_cannot_change_a_protected_input_on_any_simulation(target, operation): + simulation = built() + changed = simulation if target == "own" else built() + before = snapshot(changed) + attempts = [] + + class attempts_draw_change(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Attempt to change a stored draw" + + def formula(person, requested, parameters): + holder = changed.get_holder(DRAW) + with pytest.raises(StoredDrawError, match=DRAW): + if operation == "input": + holder.set_input(period(2026), OTHER) + elif operation == "delete": + holder.delete_arrays(period(2025)) + elif operation == "demote": + holder._memory_storage.mark_derived_except(set()) + else: + holder._memory_storage._unmark_input("default:2025") + attempts.append(operation) + return np.ones(person.count) + + simulation.tax_benefit_system.add_variable(attempts_draw_change) + for _ in range(2): + simulation.delete_arrays("attempts_draw_change", 2025) + np.testing.assert_array_equal( + simulation.calculate("attempts_draw_change", 2025), np.ones(3) + ) + assert snapshot(changed) == before + assert attempts == [operation, operation] + + +@pytest.mark.parametrize("after_read", [False, True]) +def test_an_output_callback_cannot_store_a_draw_before_or_after_a_calculation( + after_read, +): + simulation = built() + before = snapshot(simulation) + + def output_callback(current, name, requested): + if after_read: + current.calculate("age", requested) + current.get_holder(DRAW).set_input(period(2026), OTHER) + return np.ones(3) + + class draw_writing_output(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Output that attempts to change a stored draw" + calculate_output = output_callback + + simulation.tax_benefit_system.add_variable(draw_writing_output) + with pytest.raises(StoredDrawError, match=f"store an input of {DRAW}"): + simulation.calculate_output("draw_writing_output", 2025) + assert snapshot(simulation) == before + + +def test_core_restores_inputs_as_bytes_before_a_calculation_and_refuses_during_one( + tmp_path, +): + from policyengine_core.tools.simulation_dumper import _restore_holder + + directory = tmp_path / DRAW + directory.mkdir() + disk = OnDiskStorage(str(directory)) + disk.put(VALUES.copy(), period(2025)) + simulation = built({}) + + assert _restore_holder(simulation, DRAW, str(tmp_path), inputs=True) == 1 + assert_immutable(simulation.calculate(DRAW, 2025), VALUES) + before = snapshot(simulation) + with calculation(): + with pytest.raises(StoredDrawError, match=f"store an input of {DRAW}"): + _restore_holder(simulation, DRAW, str(tmp_path), inputs=True) + assert snapshot(simulation) == before diff --git a/policyengine_uk/tests/test_protected_storage.py b/policyengine_uk/tests/test_protected_storage.py index 71eaeecdc0..3ede2d86f7 100644 --- a/policyengine_uk/tests/test_protected_storage.py +++ b/policyengine_uk/tests/test_protected_storage.py @@ -49,6 +49,35 @@ def assert_immutable(value, expected): value.flat[0] = 0 +def test_an_early_year_input_delete_names_the_draw_and_branch(): + memory = ProtectedMemoryStorage(False, NAME) + early = period("0001") + memory.put(VALUES[:1], early, sequence_number=31) + before = snapshot(memory) + with calculation(): + with pytest.raises(StoredDrawError, match=f"{NAME} for 1.*default"): + memory.delete(early) + assert snapshot(memory) == before + + +@pytest.mark.parametrize("requested", [period("0001"), period(2025)]) +def test_an_early_year_delete_outside_a_calculation_keeps_cores_parser_refusal( + requested, +): + plain = InMemoryStorage(False) + memory = ProtectedMemoryStorage(False, NAME) + for target in (plain, memory): + target.put(VALUES[:1], period("0001"), sequence_number=31) + before = snapshot(memory) + with pytest.raises(ValueError) as core_error: + plain.delete(requested) + with pytest.raises(ValueError) as protected_error: + memory.delete(requested) + assert type(protected_error.value) is type(core_error.value) is ValueError + assert str(protected_error.value) == str(core_error.value) + assert snapshot(memory) == before + + @pytest.mark.parametrize("branch", ["default", "child", "ancestor"]) @pytest.mark.parametrize("target", [2025, 2026, 2027]) @pytest.mark.parametrize("values", [VALUES, OTHER]) diff --git a/policyengine_uk/tests/test_stored_draw_input_source.py b/policyengine_uk/tests/test_stored_draw_input_source.py index 9db86203a5..da4b5d843e 100644 --- a/policyengine_uk/tests/test_stored_draw_input_source.py +++ b/policyengine_uk/tests/test_stored_draw_input_source.py @@ -160,7 +160,7 @@ def test_a_model_marked_value_over_an_older_supplied_input_raises(): def test_a_value_cached_as_non_derived_over_an_input_raises(): sim = built({2025: VALUES}) sim.get_holder(DRAW).put_in_cache(OTHER, period(2025), "default", derived=False) - with pytest.raises(StoredDrawError, match="not a supplied one"): + with pytest.raises(StoredDrawError, match="holder's location record"): source_of(sim, 2025) @@ -171,7 +171,7 @@ def test_a_non_input_carry_source_raises(): np.testing.assert_array_equal(core_answer(sim, 2026), OTHER) sim._fast_cache.clear() sim.get_holder(DRAW).delete_arrays(period(2026)) - with pytest.raises(StoredDrawError, match="not a supplied input"): + with pytest.raises(StoredDrawError, match="simulation's input record"): sim.calculate(DRAW, 2026) @@ -710,10 +710,33 @@ def test_a_source_replaced_while_core_reads_it_raises(monkeypatch): def replaces_the_source(self, name, *args, **kwargs): result = real(self, name, *args, **kwargs) if name == DRAW: - # The same values, stored again: a new sequence number. + # The holder must refuse even an equal replacement. self.get_holder(DRAW).set_input(period(2025), VALUES) return result monkeypatch.setattr(CoreSimulation, "calculate", replaces_the_source) + source = source_of(sim, 2025) + input_keys = set(sim._user_input_keys) + locations = dict(sim.get_holder(DRAW)._user_input_storage) + with pytest.raises(StoredDrawError, match="during a calculation"): + sim.calculate(DRAW, 2025) + assert source_of(sim, 2025) == source + assert sim._user_input_keys == input_keys + assert sim.get_holder(DRAW)._user_input_storage == locations + np.testing.assert_array_equal(sim.get_holder(DRAW).get_array(period(2025)), VALUES) + + +def test_the_answer_check_refuses_a_changed_source_sequence(monkeypatch): + sim = built({2025: VALUES}) + real = CoreSimulation.calculate + + def changes_the_sequence(self, name, *args, **kwargs): + result = real(self, name, *args, **kwargs) + if name == DRAW: + memory = self.get_holder(DRAW)._memory_storage + memory._sequence_numbers["default:2025"] += 1 + return result + + monkeypatch.setattr(CoreSimulation, "calculate", changes_the_sequence) with pytest.raises(StoredDrawError, match="changed while"): sim.calculate(DRAW, 2025) diff --git a/policyengine_uk/utils/protected_inputs.py b/policyengine_uk/utils/protected_inputs.py new file mode 100644 index 0000000000..193e2bec76 --- /dev/null +++ b/policyengine_uk/utils/protected_inputs.py @@ -0,0 +1,86 @@ +"""Install protected draw storage and check its input records before reads.""" + +from policyengine_core.data_storage import InMemoryStorage +from policyengine_core.data_storage.store_history import _period_from_storage_key + +from policyengine_uk.utils.data_source import built_from_data +from policyengine_uk.utils.protected_storage import ProtectedMemoryStorage +from policyengine_uk.utils.stored_draws import REGISTERED_DRAW_NAMES, StoredDrawError + + +def protects_draws(simulation) -> bool: + """Whether the current population attachment installed draw protection.""" + return bool(getattr(simulation, "_stored_draws_protected", False)) + + +def install_draw_storage(simulation) -> None: + """Snapshot registered holders before a data loader supplies its values.""" + protected = built_from_data(simulation) + simulation._stored_draws_protected = protected + for name in REGISTERED_DRAW_NAMES: + if name not in simulation.tax_benefit_system.variables: + continue + if protected: + holder = simulation.get_holder(name) + else: + entity = simulation.tax_benefit_system.variables[name].entity.key + holder = simulation.populations[entity]._holders.get(name) + if holder is None: + continue + storage = holder._memory_storage + if protected: + if not isinstance(storage, ProtectedMemoryStorage): + holder._memory_storage = ProtectedMemoryStorage.from_storage( + storage, name + ) + elif isinstance(storage, ProtectedMemoryStorage): + # A situation owns ordinary writeable values, including adopted inputs. + state = storage.__getstate__() + state.pop("_array_dtypes", None) + state.pop("variable_name", None) + state["_arrays"] = { + key: values.copy() for key, values in state["_arrays"].items() + } + plain = InMemoryStorage(storage.is_eternal) + plain.__setstate__(state) + holder._memory_storage = plain + + +def check_draw_storage(simulation, holder, name: str) -> None: + """Require protected memory and matching provenance for every visible mark.""" + storage = holder._memory_storage + if not isinstance(storage, ProtectedMemoryStorage): + raise StoredDrawError(f"{name} needs protected memory storage for every read.") + visible = set(simulation._get_visible_branch_names()) + recorded = getattr(simulation, "_user_input_keys", None) + for storage_key in storage._inputs: + if storage_key not in storage._arrays: + continue + branch_name, period_string = storage_key.split(":", 1) + if branch_name not in visible: + continue + stored_period = _period_from_storage_key(period_string) + stored_period = holder._storage_period(stored_period) + key = (name, branch_name, stored_period) + description = f"{name} for {stored_period} on branch {branch_name!r}" + if recorded is None or key not in recorded: + raise StoredDrawError( + f"Protected storage holds {description} as an input, but the " + "simulation's input record does not name it as a supplied input." + ) + # Core seeds register-only inputs before testing their storage locations. + holder._seed_input_storage(stored_period, branch_name) + location = getattr(holder, "_user_input_storage", {}).get( + (branch_name, str(stored_period)) + ) + if ( + location is None + or location[0] != stored_period + or "memory" not in location[1] + or not holder._stores_user_input(stored_period, branch_name, "memory") + ): + raise StoredDrawError( + f"Protected storage holds {description} as an input, but the " + "holder's location record does not name it as a supplied input " + "in memory." + ) diff --git a/policyengine_uk/utils/protected_storage.py b/policyengine_uk/utils/protected_storage.py index 3a4ea0d5a2..6a89abc48b 100644 --- a/policyengine_uk/utils/protected_storage.py +++ b/policyengine_uk/utils/protected_storage.py @@ -3,6 +3,7 @@ import numpy as np from policyengine_core import periods from policyengine_core.data_storage import InMemoryStorage +from policyengine_core.data_storage.store_history import _period_from_storage_key from policyengine_uk.utils.stored_draws import ( StoredDrawError, @@ -97,7 +98,8 @@ def delete(self, period=None, branch_name="default"): for key in self._arrays if key.startswith(prefix) and ( - period is None or period.contains(periods.period(key.split(":", 1)[1])) + period is None + or period.contains(_period_from_storage_key(key.split(":", 1)[1])) ) } selected_inputs = selected.intersection(self._inputs) diff --git a/policyengine_uk/utils/stored_draws.py b/policyengine_uk/utils/stored_draws.py index 534a54a92d..d90507c2b4 100644 --- a/policyengine_uk/utils/stored_draws.py +++ b/policyengine_uk/utils/stored_draws.py @@ -268,6 +268,13 @@ def input_source(population, name: str, period) -> Optional[InputSource]: f"with an underscore ({CORE_DISK_KEY_ISSUE}). Read {name} from a " "simulation without memory_config." ) + from policyengine_uk.utils.protected_inputs import ( + check_draw_storage, + protects_draws, + ) + + if protects_draws(simulation): + check_draw_storage(simulation, holder, name) storing = holder._branch_storing(period, branch_name) if storing is not None: memory = holder._memory_storage From 10b176a1b419481ed63b1022650bf3b296cfe5d8 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 8 Nov 2026 12:04:34 -0500 Subject: [PATCH 11/11] Adoption tests mutate the donor input through its owner Co-Authored-By: Claude Opus 5.5 --- policyengine_uk/tests/test_protected_draw_storage.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/policyengine_uk/tests/test_protected_draw_storage.py b/policyengine_uk/tests/test_protected_draw_storage.py index c8377b6e35..98b13bd004 100644 --- a/policyengine_uk/tests/test_protected_draw_storage.py +++ b/policyengine_uk/tests/test_protected_draw_storage.py @@ -167,6 +167,9 @@ def test_root_clone_and_branch_keep_input_marks_values_and_sequence_numbers(): def test_a_data_rebuild_adopts_plain_inputs_without_changing_their_provenance(): source = Simulation(situation=situation()) source_holder = source.get_holder(DRAW) + source_values = VALUES.copy() + source_holder.set_input(period(2025), source_values) + assert source_holder.get_array(period(2025)) is source_values before = snapshot(source) simulation = built() simulation.build_from_populations(source.populations) @@ -174,7 +177,7 @@ def test_a_data_rebuild_adopts_plain_inputs_without_changing_their_provenance(): assert snapshot(simulation) == before assert source_of(simulation, 2025).sequence_number == before[2]["default:2025"] assert_immutable(simulation.calculate(DRAW, 2025), VALUES) - source_holder.get_array(period(2025))[0] = 0 + source_values[0] = 0 assert_immutable(simulation.calculate(DRAW, 2025), VALUES) assert simulation.populations["person"] is not source.populations["person"] assert simulation.get_holder(DRAW) is not source_holder