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/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/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/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/simulation.py b/policyengine_uk/simulation.py index e000f0f9a3..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, @@ -99,6 +103,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 +198,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 +653,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 @@ -1423,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( @@ -1444,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]: @@ -1799,6 +1832,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 +1854,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 +1878,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/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_no_engine_time_randomness.py b/policyengine_uk/tests/code_health/test_no_engine_time_randomness.py new file mode 100644 index 0000000000..6c60b15e90 --- /dev/null +++ b/policyengine_uk/tests/code_health/test_no_engine_time_randomness.py @@ -0,0 +1,4210 @@ +"""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. 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 +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. 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 +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 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, + }, + # 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: 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, + }, + # 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()), + # 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 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) + ) 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/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_stored_draw_consumers.py b/policyengine_uk/tests/test_stored_draw_consumers.py new file mode 100644 index 0000000000..19f58e3f65 --- /dev/null +++ b/policyengine_uk/tests/test_stored_draw_consumers.py @@ -0,0 +1,500 @@ +"""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 + +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" +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 + + +@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 + ) 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 new file mode 100644 index 0000000000..9db86203a5 --- /dev/null +++ b/policyengine_uk/tests/test_stored_draw_input_source.py @@ -0,0 +1,719 @@ +"""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, +) + +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(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}, + {"requires_computation_after": "age"}, + ): + 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_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( + 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]) + ) + + +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] + ) + + +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/tests/test_supplied_input_tightening.py b/policyengine_uk/tests/test_supplied_input_tightening.py new file mode 100644 index 0000000000..eaab3a244d --- /dev/null +++ b/policyengine_uk/tests/test_supplied_input_tightening.py @@ -0,0 +1,205 @@ +"""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, +) + +pytestmark = pytest.mark.usefixtures("cloned_uk_tax_benefit_system") + +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/stored_draws.py b/policyengine_uk/utils/stored_draws.py new file mode 100644 index 0000000000..534a54a92d --- /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: + 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 + 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) + _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) 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