Coverage for src/basanos/math/_stream_io.py: 100%

1"""Persistence for `BasanosStream` — save/load of the full stream state. 

2 

3The archive layout is versioned by `_SAVE_FORMAT_VERSION`. `load_stream_archive` 

4validates the version and the presence of every required key before 

5deserialising anything, so a stale or hand-edited archive fails with a 

6descriptive error rather than a bare ``KeyError``. 

7 

8These are free functions (rather than `BasanosStream` methods) so the 

9serialisation format depends only on the state layout in 

10:mod:`basanos.math._stream_state`, not on the stream façade. 

11""" 

12 

13from __future__ import annotations 

14 

15import dataclasses 

16import os 

17from typing import Any 

18 

19import numpy as np 

20 

21from ..exceptions import StreamStateCorruptError 

22from ._config import BasanosConfig 

23from ._stream_state import _REQUIRED_KEYS, _SAVE_FORMAT_VERSION, _StreamState 

24 

25 

26def save_stream_archive( 

27 cfg: BasanosConfig, 

28 assets: list[str], 

29 state: _StreamState, 

30 path: str | os.PathLike[str], 

31) -> None: 

32 """Serialise a stream's config, assets, and state to a ``.npz`` archive. 

33 

34 All `_StreamState` arrays, the configuration, and the asset 

35 list are written in a single `np.savez` call. A stream 

36 restored via `load_stream_archive` produces bit-for-bit identical 

37 `BasanosStream.step` output. 

38 

39 Args: 

40 cfg: The stream configuration to serialise. 

41 assets: Ordered asset column names. 

42 state: The mutable stream state carrier to serialise. 

43 path: Destination file path. `np.savez` appends 

44 ``.npz`` automatically when the suffix is absent. 

45 """ 

46 # Build the per-field dict automatically from _StreamState so that any 

47 # new field added to the dataclass is included without manual updates. 

48 state_arrays: dict[str, Any] = {} 

49 for field in dataclasses.fields(_StreamState): 

50 value = getattr(state, field.name) 

51 if field.name in ("sw_ret_buf", "corr_ret_buf"): 

52 # Sentinel: use an empty (0, 0) array to represent None so the 

53 # key is always present in the archive and load() can detect it. 

54 state_arrays[field.name] = value if value is not None else np.empty((0, 0), dtype=float) 

55 elif field.name == "step_count": 

56 state_arrays[field.name] = np.array(value) 

57 else: 

58 state_arrays[field.name] = value 

59 np.savez( 

60 path, 

61 format_version=np.array(_SAVE_FORMAT_VERSION), 

62 cfg_json=np.array(cfg.model_dump_json()), 

63 assets=np.array(assets), 

64 **state_arrays, 

65 ) 

66 

67 

68def _validate_archive_version(data: Any) -> None: 

69 """Raise if *data* lacks a format tag or has an incompatible version.""" 

70 if "format_version" not in data: 

71 raise ValueError( # noqa: TRY003 

72 "Stream file is missing a format version tag. " 

73 "It was written with an incompatible version of BasanosStream. " 

74 "Re-generate it via BasanosStream.from_warmup()." 

75 ) 

76 found = int(data["format_version"]) 

77 if found != _SAVE_FORMAT_VERSION: 

78 raise ValueError( # noqa: TRY003 

79 f"Stream file was written with format version {found}, " 

80 f"but the current version is {_SAVE_FORMAT_VERSION}. " 

81 "Re-generate it via BasanosStream.from_warmup()." 

82 ) 

83 

84 

85def _state_field_from_archive(field_name: str, raw: Any) -> Any: 

86 """Deserialise a single `_StreamState` field from its archive representation.""" 

87 if field_name in ("sw_ret_buf", "corr_ret_buf"): 

88 return raw if raw.size > 0 else None 

89 if field_name == "step_count": 

90 return int(raw) 

91 return raw 

92 

93 

94def load_stream_archive( 

95 path: str | os.PathLike[str], 

96) -> tuple[BasanosConfig, list[str], _StreamState]: 

97 """Restore a stream's ``(cfg, assets, state)`` from a saved ``.npz`` archive. 

98 

99 Args: 

100 path: Path to a ``.npz`` archive written by `save_stream_archive`. 

101 

102 Returns: 

103 A ``(cfg, assets, state)`` tuple whose reconstructed state reproduces 

104 the original stream bit-for-bit at the time it was saved. 

105 

106 Raises: 

107 ValueError: If the archive is missing its format-version tag or was 

108 written with an incompatible format version. 

109 StreamStateCorruptError: If a required key is absent from the archive. 

110 """ 

111 with np.load(path, allow_pickle=False) as data: 

112 _validate_archive_version(data) 

113 # Validate that every required key is present. This catches archives 

114 # that were produced by an older codebase missing a newly added field, 

115 # or archives that have been manually edited, with a descriptive error 

116 # instead of a bare KeyError. 

117 archive_keys = frozenset(data.files) 

118 missing = _REQUIRED_KEYS - archive_keys 

119 if missing: 

120 raise StreamStateCorruptError(missing) 

121 cfg = BasanosConfig.model_validate_json(data["cfg_json"].item()) 

122 assets: list[str] = list(data["assets"]) 

123 state_kwargs: dict[str, Any] = { 

124 field.name: _state_field_from_archive(field.name, data[field.name]) 

125 for field in dataclasses.fields(_StreamState) 

126 } 

127 state = _StreamState(**state_kwargs) 

128 return cfg, assets, state