Keep branch period reads, dumps, and disk deletion consistent

changelog.d/fix-uprating-branch-visibility.fixed.md

@@ -0,0 +1 @@
+Scope known periods and dumps to the branch's readable values while preserving derived marks, parse underscore-containing disk branch names safely, delete contained disk periods consistently with memory, reject ambiguous or unparseable storage keys before writing, and warn and skip malformed restored filenames.

docs/python_api/data_storage.md

@@ -20,6 +20,20 @@ The `policyengine_core.data_storage` module contains the classes that are used t
     :show-inheritance:
 ```
 
+## Storage keys
+
+Both storage backends reject branch names containing `:` and month or year
+periods anchored after the first day of a month. They also reject periods whose
+string form cannot be parsed back, such as the year `0999`, before storing a
+value. Existing aliases, including twelve-month periods serialized as one year,
+remain supported. Disk restoration warns and skips filenames without a branch
+separator and a canonical period string.
+
+```{eval-rst}
+.. automodule:: policyengine_core.data_storage.storage_keys
+    :members:
+```
+
 ## TemporaryStorageDirectory
 
 ```{eval-rst}

policyengine_core/data_storage/in_memory_storage.py

@@ -4,6 +4,7 @@
 from numpy.typing import ArrayLike
 
 from policyengine_core import periods
+from policyengine_core.data_storage.storage_keys import check_storable
 from policyengine_core.periods import Period
 
 
@@ -266,20 +267,8 @@ def put(
         period = periods.period(period)
 
         # Keys embed f"{branch_name}:{period}", so reject inputs whose key
-        # would be ambiguous to parse back: branch names containing the
-        # separator, and month/year periods anchored mid-month (their string
-        # form drops the day). See policyengine-core#526 for the structured-
-        # key refactor that will lift these restrictions.
-        if ":" in branch_name:
-            raise ValueError(
-                f"Branch name {branch_name!r} may not contain ':' "
-                "(see policyengine-core#526)."
-            )
-        if period.unit in (periods.YEAR, periods.MONTH) and period.start.day != 1:
-            raise ValueError(
-                f"Cannot cache period {period} anchored mid-month: its "
-                "string form is lossy (see policyengine-core#526)."
-            )
+        # would be ambiguous to parse back, as ``OnDiskStorage.put`` does.
+        check_storable(branch_name, period)
         key = f"{branch_name}:{period}"
         self._arrays[key] = value
         self._stop_sharing(key)

policyengine_core/data_storage/on_disk_storage.py

@@ -1,6 +1,7 @@
 import os
 import shutil
 import tempfile
+import warnings
 import weakref
 
 import numpy
@@ -11,6 +12,7 @@
     directory_containing,
     path_key,
 )
+from policyengine_core.data_storage.storage_keys import check_storable
 from policyengine_core.enums import EnumArray
 from policyengine_core.periods import Period
 
@@ -48,6 +50,41 @@ def _epoch() -> tuple:
     return os.getpid(), _forks
 
 
+def _split_key(key: str) -> tuple:
+    """Split a ``f"{branch_name}_{period}"`` file key into its two parts.
+
+    Branch names often contain ``_`` (policyengine-us uses ``no_salt`` and
+    ``mtr_for_adult_1``) but a period's string form never does, so the key
+    splits on its last ``_``.
+    """
+    branch_name, period = key.rsplit("_", 1)
+    return branch_name, period
+
+
+def _is_key(key: str) -> bool:
+    """Whether ``key`` ends with a canonical period after the branch separator.
+
+    Legacy branch names and eternal-storage files retain their restore
+    behavior; this check only validates the period's spelling.
+    """
+    _, separator, period = key.rpartition("_")
+    if not separator:
+        return False
+    try:
+        return str(periods.period(period)) == period
+    except ValueError:
+        return False
+
+
+def _is_within(key: str, branch_name: str, period: Period) -> bool:
+    """Whether file key ``key`` stores a value of ``branch_name`` for a
+    period that ``period`` contains."""
+    key_branch_name, key_period = _split_key(key)
+    return key_branch_name == branch_name and period.contains(
+        periods.period(key_period)
+    )
+
+
 class OnDiskStorage:
     """
     Low-level class responsible for storing and retrieving calculated vectors on disk
@@ -275,6 +312,11 @@ def put(
         if self.is_eternal:
             period = periods.period(periods.ETERNITY)
         period = periods.period(period)
+        # The key must split back into this branch name and period (see
+        # ``get_known_branch_periods`` and ``delete``): reject what
+        # ``InMemoryStorage.put`` rejects. A month anchored mid-month, say,
+        # would be stored under the key of the calendar month.
+        check_storable(branch_name, period)
 
         filename = f"{branch_name}_{period}"
         path = self._path_to_write(filename)
@@ -328,12 +370,13 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None:
         if period is None:
             # Only wipe files belonging to the requested branch (previously
             # this wiped every branch regardless of ``branch_name`` — same
-            # class of bug as C2 in InMemoryStorage).
-            branch_prefix = f"{branch_name}_"
+            # class of bug as C2 in InMemoryStorage). Compare the parsed
+            # branch name, not a prefix: deleting ``pre_tcja`` must not
+            # also wipe ``pre_tcja_ctc``.
             self._files = {
                 period_item: value
                 for period_item, value in self._files.items()
-                if not period_item.startswith(branch_prefix)
+                if _split_key(period_item)[0] != branch_name
             }
             self._derived.intersection_update(self._files)
             return
@@ -342,27 +385,34 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None:
             period = periods.period(periods.ETERNITY)
         period = periods.period(period)
 
+        # Delete every period of the branch that ``period`` contains, as
+        # ``InMemoryStorage.delete`` does and ``Holder.delete_arrays``
+        # documents: deleting ``2025`` also deletes ``2025-01``. Previously
+        # only the file keyed exactly ``period`` was deleted.
         if period is not None:
             self._files = {
                 period_item: value
                 for period_item, value in self._files.items()
-                if not period_item == f"{branch_name}_{period}"
+                if not _is_within(period_item, branch_name, period)
             }
             self._derived.intersection_update(self._files)
 
     def get_known_periods(self) -> list:
-        return list([periods.period(x.split("_")[1]) for x in self._files.keys()])
+        return [period for _, period in self.get_known_branch_periods()]
 
     def get_known_branch_periods(self) -> list:
         return [
             (branch_name, periods.period(period))
-            for branch_name, period in map(lambda x: x.split("_"), self._files.keys())
+            for branch_name, period in map(_split_key, self._files.keys())
         ]
 
     def restore(self) -> None:
         """Read back the values stored in this storage's directory: for each
         key, the file named for it (not those in ``REPLACEMENTS_DIR``).
 
+        A ``.npy`` file whose name lacks a branch separator and a canonical
+        period string is left out, with a warning.
+
         This storage writes over those files, as over those it wrote, until
         it is cloned or copied, or this process forks: from then on it writes
         a new file instead, since the clone, copy or forked process reads
@@ -374,14 +424,28 @@ def restore(self) -> None:
         self._files = files = {}
         # Files read back from a directory carry no derived marks.
         self._derived = set()
+        not_keys = []
         # Restore self._files from content of storage_dir.
         for filename in os.listdir(self.storage_dir):
             if not filename.endswith(".npy"):
                 continue
             path = os.path.join(self.storage_dir, filename)
             filename_core = filename.rsplit(".", 1)[0]
+            if not _is_key(filename_core):
+                # Reading it back would put a key in ``_files`` that names no
+                # branch and period, which ``delete`` and
+                # ``get_known_periods`` could not parse.
+                not_keys.append(filename)
+                continue
             files[filename_core] = path
         self._restored_paths.update(files.values())
+        if not_keys:
+            warnings.warn(
+                f"Not restoring {', '.join(sorted(not_keys))} in "
+                f"{self.storage_dir}: a stored value's file is named for a "
+                "branch name, '_' and a period.",
+                stacklevel=2,
+            )
 
     def __del__(self, _rmtree=shutil.rmtree, _getpid=os.getpid) -> None:
         # (The defaults keep the two functions reachable while the

policyengine_core/data_storage/storage_keys.py

@@ -0,0 +1,43 @@
+"""What a storage key can hold.
+
+Both storages key a value by its branch name and its period's string form:
+``f"{branch_name}:{period}"`` in memory, ``f"{branch_name}_{period}"`` on disk
+(the name of the value's file). ``check_storable`` enforces the same branch
+and period serialization restrictions in both storages (``Holder._set``
+chooses between them by how much memory is in use).
+"""
+
+from policyengine_core import periods
+from policyengine_core.periods import Period
+
+
+def check_storable(branch_name: str, period: Period) -> None:
+    """Raise ``ValueError`` for an ambiguous branch or period serialization.
+
+    Reject a branch name containing ``:``, which separates the parts of an
+    in-memory key, and a year or month anchored mid-month, whose string form
+    drops the day (``month:2025-03-15`` is ``2025-03``), so that its key would
+    be the key of another period. Reject unparseable period strings too,
+    including years below 1000 whose string form loses leading zeroes.
+    Supported aliases, such as twelve months serialized as one year, retain
+    their existing meaning. See policyengine-core#526 for the structured-key
+    refactor.
+    """
+    if ":" in branch_name:
+        raise ValueError(
+            f"Branch name {branch_name!r} may not contain ':' "
+            "(see policyengine-core#526)."
+        )
+    if period.unit in (periods.YEAR, periods.MONTH) and period.start.day != 1:
+        raise ValueError(
+            f"Cannot cache period {period} anchored mid-month (it starts on "
+            f"{period.start}): its string form is lossy (see "
+            "policyengine-core#526)."
+        )
+    try:
+        periods.period(str(period))
+    except ValueError as error:
+        raise ValueError(
+            f"Cannot cache period {period!r}: its string form cannot be parsed "
+            "(see policyengine-core#526)."
+        ) from error

policyengine_core/holders/holder.py

@@ -205,11 +205,25 @@ def get_memory_usage(self) -> dict:
 
         return usage
 
-    def get_known_periods(self) -> List[Period]:
+    def get_known_periods(self, branch_name: str = None) -> List[Period]:
         """
         Get the list of periods the variable value is known for.
-        """
 
+        With ``branch_name``, list each period ``get_array(period,
+        branch_name)`` reads a stored value for, once: those stored under
+        that branch, its ``parent_branch`` ancestors or ``default`` (see
+        ``_readable_branches``). Without it, list the period of every value
+        stored, under any branch; that branch may not read some of them.
+        """
+        if branch_name is not None:
+            readable = set(self._readable_branches(branch_name))
+            return list(
+                dict.fromkeys(
+                    period
+                    for stored_branch, period in self.get_known_branch_periods()
+                    if stored_branch in readable
+                )
+            )
         return list(self._memory_storage.get_known_periods()) + list(
             (self._disk_storage.get_known_periods() if self._disk_storage else [])
         )

policyengine_core/simulations/simulation.py

@@ -1082,8 +1082,12 @@ def _calculate(self, variable_name: str, period: Period = None) -> ArrayLike:
 
             # If no result, use the default value and cache it
             if array is None:
-                # Check if the variable has a previously defined value
-                known_periods = holder.get_known_periods()
+                # Check if the variable has a previously defined value.
+                # Only periods this branch can read count: a value stored
+                # only under a branch this one cannot read must not change
+                # what it calculates or caches (below, a later stored period
+                # leaves the default uncached).
+                known_periods = holder.get_known_periods(self.branch_name)
                 earlier_known_periods = [
                     known_period
                     for known_period in known_periods

policyengine_core/tools/simulation_dumper.py

@@ -70,13 +70,17 @@ def restore_simulation(directory, tax_benefit_system, **kwargs):
 
 def _dump_holder(holder, directory):
     disk_storage = holder.create_disk_storage(directory, preserve=True)
+    # Dump what the simulation reads: on a branch, its own value, else its
+    # nearest ancestor's or the default one. A value stored only under a
+    # branch the simulation cannot read is left out: read here, it would be
+    # ``None``, saved as an array ``restore_simulation`` cannot load.
+    branch_name = holder.simulation.branch_name
     derived_periods = set()
-    for period in holder.get_known_periods():
-        value = holder.get_array(period)
-        disk_storage.put(value, period)
+    for period in holder.get_known_periods(branch_name):
+        disk_storage.put(holder.get_array(period, branch_name), period)
         # Read the mark of exactly the value dumped: the same period on the
         # same branch as ``get_array``.
-        if holder.is_derived(period):
+        if holder.is_derived(period, branch_name):
             derived_periods.add(str(period))
     if derived_periods:
         path = os.path.join(disk_storage.storage_dir, DERIVED_PERIODS_FILE)