diff --git a/nslsii/ophyd_async/providers.py b/nslsii/ophyd_async/providers.py index 92ac0d20..673a84f2 100644 --- a/nslsii/ophyd_async/providers.py +++ b/nslsii/ophyd_async/providers.py @@ -140,7 +140,7 @@ class NSLS2PathProvider(PathProvider): /nsls2/data/{TLA}/proposals/{CYCLE}/{PROPOSAL}/assets/{DETECTOR}/{Y}/{M}/{D} If include_scan_id_dir is provided, an additional directory will be created at the - end of the path with the name scan_{scan_id:06}. + end of the path with the name scan_{scan_id} zero-padded to scan_id_padding digits. Parameters ---------- @@ -164,6 +164,8 @@ class NSLS2PathProvider(PathProvider): If set to YMDGranularity.year, the path will include only the year directory. include_scan_id_dir : bool, default False Whether to include a scan ID directory at the end of the path. + scan_id_padding : int, default 6 + Number of digits to zero-pad the scan ID to when generating the scan ID directory name. """ def __init__( @@ -175,6 +177,7 @@ def __init__( granularity: YMDGranularity | str = YMDGranularity.day, separator: str | None = None, include_scan_id_dir: bool = False, + scan_id_padding: int = 6, ): self._filename_provider = filename_provider @@ -199,6 +202,9 @@ def __init__( self._ymd_separator = separator or ("\\" if isinstance(self._base_write_dir, PureWindowsPath) else "/") self._include_scan_id_dir = include_scan_id_dir + if isinstance(scan_id_padding, bool) or not isinstance(scan_id_padding, int) or scan_id_padding < 0: + raise ValueError("scan_id_padding must be a non-negative integer!") + self._scan_id_padding = scan_id_padding @property def filename_provider(self) -> FilenameProvider: @@ -211,7 +217,7 @@ def _generate_directory_path(self, base_path: PurePath, datakey_name: str) -> Pu Depending on the granularity, the path will include year, month, and day directories. If the include_scan_id_dir is provided, an additional directory will be created at the end of the - path with the name scan_{scan_id:06}. + path with the name scan_{scan_id} zero-padded to scan_id_padding digits. Parameters ---------- @@ -254,7 +260,10 @@ def _generate_directory_path(self, base_path: PurePath, datakey_name: str) -> Pu ) if self._include_scan_id_dir: - directory_path = directory_path / f"scan_{self._metadata_dict['scan_id']:06}" + directory_path = ( + directory_path + / f"scan_{self._metadata_dict['scan_id']:0{self._scan_id_padding}}" + ) return directory_path diff --git a/nslsii/tests/ophyd_async/test_ophyd_async_providers.py b/nslsii/tests/ophyd_async/test_ophyd_async_providers.py index ed73bb66..c521e745 100644 --- a/nslsii/tests/ophyd_async/test_ophyd_async_providers.py +++ b/nslsii/tests/ophyd_async/test_ophyd_async_providers.py @@ -135,6 +135,36 @@ def test_nsls2_path_provider_requires_tla(dummy_re_md_dict, monkeypatch): NSLS2PathProvider(dummy_re_md_dict) +@pytest.mark.parametrize( + ("scan_id_padding", "expected_scan_dir"), + [ + (6, "scan_000005"), + (0, "scan_5"), + (3, "scan_005"), + (10, "scan_0000000005"), + ], +) +def test_nsls2_path_provider_scan_id_padding( + dummy_re_md_dict, static_fp, scan_id_padding, expected_scan_dir +): + os.environ["BEAMLINE_ACRONYM"] = "tst" + pp = NSLS2PathProvider( + dummy_re_md_dict, + filename_provider=static_fp, + include_scan_id_dir=True, + scan_id_padding=scan_id_padding, + ) + info = pp("test") + assert str(info.directory_path).endswith(expected_scan_dir) + + +@pytest.mark.parametrize("scan_id_padding", [-1, 1.5, True, "6"]) +def test_nsls2_path_provider_rejects_invalid_padding(dummy_re_md_dict, scan_id_padding): + os.environ["BEAMLINE_ACRONYM"] = "tst" + with pytest.raises(ValueError, match="scan_id_padding must be a non-negative integer"): + NSLS2PathProvider(dummy_re_md_dict, scan_id_padding=scan_id_padding) + + FIXED_NOW = datetime(2025, 3, 15, 10, 30, 45)