diff --git a/runs/fidelity-headroom/prepare_pilot_simulator.py b/runs/fidelity-headroom/prepare_pilot_simulator.py index c3d77fb..fc84bde 100644 --- a/runs/fidelity-headroom/prepare_pilot_simulator.py +++ b/runs/fidelity-headroom/prepare_pilot_simulator.py @@ -32,6 +32,9 @@ PRIMARY_ROLES = ("low1", "high1") def load_module(path: Path): + module_root = str(path.parent.resolve()) + if module_root not in sys.path: + sys.path.insert(0, module_root) spec = importlib.util.spec_from_file_location("simfid_s2rb_prepare", path) if spec is None or spec.loader is None: raise ImportError(path) diff --git a/runs/fidelity-headroom/run_pilot_simulator.py b/runs/fidelity-headroom/run_pilot_simulator.py index d393916..6db18b4 100644 --- a/runs/fidelity-headroom/run_pilot_simulator.py +++ b/runs/fidelity-headroom/run_pilot_simulator.py @@ -16,6 +16,9 @@ from typing import Any def load_module(name: str, path: Path): + module_root = str(path.parent.resolve()) + if module_root not in sys.path: + sys.path.insert(0, module_root) spec = importlib.util.spec_from_file_location(name, path) if spec is None or spec.loader is None: raise ImportError(path) diff --git a/runs/fidelity-headroom/test_strong_pilot.py b/runs/fidelity-headroom/test_strong_pilot.py index 48f6089..921c2b3 100644 --- a/runs/fidelity-headroom/test_strong_pilot.py +++ b/runs/fidelity-headroom/test_strong_pilot.py @@ -8,6 +8,8 @@ from pathlib import Path import numpy as np from analyze_prefixes import PrefixExample +from prepare_pilot_simulator import load_module as load_prepare_module +from run_pilot_simulator import load_module as load_run_module from analyze_strong_pilot import ( fit_model, load_pilot_simulator, @@ -68,6 +70,19 @@ def main() -> None: features, red_flags = load_pilot_simulator(path) assert len(features) == 12 assert red_flags == [] + + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + (root / "prepare_dependency.py").write_text("VALUE = 17\n", encoding="utf-8") + (root / "prepare_target.py").write_text( + "from prepare_dependency import VALUE\n", encoding="utf-8" + ) + assert load_prepare_module(root / "prepare_target.py").VALUE == 17 + (root / "run_dependency.py").write_text("VALUE = 23\n", encoding="utf-8") + (root / "run_target.py").write_text( + "from run_dependency import VALUE\n", encoding="utf-8" + ) + assert load_run_module("run_target", root / "run_target.py").VALUE == 23 print("fidelity strong pilot: PASS")