|
4 | 4 | """ |
5 | 5 |
|
6 | 6 | from dataclasses import dataclass, field |
| 7 | +from importlib.resources import as_file, files |
7 | 8 | from pathlib import Path |
8 | 9 | import sys |
9 | 10 | import tomllib |
10 | 11 | from typing import Literal |
11 | 12 |
|
12 | 13 | REPO_ROOT = Path(__file__).parent.parent.parent |
13 | | -REGIMES_PATH = REPO_ROOT / "data" / "regimes.toml" |
14 | | -TRANSFORMATIONS_PATH = REPO_ROOT / "data" / "transformations.toml" |
| 14 | +DATA_ROOT = files("se_regimes") / "data" |
| 15 | +REGIMES_RESOURCE = DATA_ROOT / "regimes.toml" |
| 16 | +TRANSFORMATIONS_RESOURCE = DATA_ROOT / "transformations.toml" |
15 | 17 |
|
16 | 18 | Outcome = Literal["PRS", "BRK", "INH", "TODO", "—"] |
17 | 19 |
|
@@ -122,13 +124,26 @@ def coverage_matrix(self) -> dict[str, dict[str, Outcome]]: |
122 | 124 |
|
123 | 125 |
|
124 | 126 | def load( |
125 | | - regimes_path: Path = REGIMES_PATH, |
126 | | - transformations_path: Path = TRANSFORMATIONS_PATH, |
| 127 | + regimes_path: Path | None = None, |
| 128 | + transformations_path: Path | None = None, |
127 | 129 | ) -> Registry: |
128 | | - """Load the registry from the given TOML files, returning a Registry object.""" |
129 | | - with Path.open(regimes_path, "rb") as f: |
| 130 | + """Load the registry from TOML files, returning a Registry object.""" |
| 131 | + if regimes_path is not None and transformations_path is not None: |
| 132 | + return _load_from_paths(regimes_path, transformations_path) |
| 133 | + |
| 134 | + with ( |
| 135 | + as_file(REGIMES_RESOURCE) as packaged_regimes_path, |
| 136 | + as_file(TRANSFORMATIONS_RESOURCE) as packaged_transformations_path, |
| 137 | + ): |
| 138 | + return _load_from_paths(packaged_regimes_path, packaged_transformations_path) |
| 139 | + |
| 140 | + |
| 141 | +def _load_from_paths(regimes_path: Path, transformations_path: Path) -> Registry: |
| 142 | + """Load the registry from concrete filesystem paths.""" |
| 143 | + with regimes_path.open("rb") as f: |
130 | 144 | raw_regimes = tomllib.load(f) |
131 | | - with Path.open(transformations_path, "rb") as f: |
| 145 | + |
| 146 | + with transformations_path.open("rb") as f: |
132 | 147 | raw_families = tomllib.load(f) |
133 | 148 |
|
134 | 149 | registry = Registry() |
|
0 commit comments