Add scan_id_padding param to NSLS2PathProvider

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
 

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)