Repository: bluesky/bluesky
Pull request #1808 (2024-10-08 23:12:48)

Determine indices of tiff files from StreamDatum documents
Uses the `indices` field of StreamDatum documents to set the indices in the templated file names of tiff files in tiff collection. Assumes that each file contains `chunk_shape[0]` frames and the indices of files are consecutive.

Tests: TBD


Patch:
diff --git a/src/bluesky/consolidators.py b/src/bluesky/consolidators.py
index 265318dc..cf9e7980 100644
--- a/src/bluesky/consolidators.py
+++ b/src/bluesky/consolidators.py
@@ -288,16 +288,19 @@ class TIFFConsolidator(ConsolidatorBase):
         return self.uri + self._sres_parameters["template"].format(indx)
 
     def consume_stream_datum(self, doc: StreamDatum):
-        indx = int(doc["uid"].split("/")[1])
-        new_datum_uri = self.get_datum_uri(indx)
-        new_asset = Asset(
-            data_uri=new_datum_uri,
-            is_directory=False,
-            parameter="data_uris",
-            num=len(self.assets) + 1,
-        )
-        self.assets.append(new_asset)
-        self.data_uris.append(new_datum_uri)
+        # Determine the indices in the names of tiff files from indices of frames and number of frames per file
+        first_file_indx = int(doc["indices"]["start"] / self.chunk_shape[0])
+        last_file_indx = int((doc["indices"]["stop"] - 1) / self.chunk_shape[0])
+        for indx in range(first_file_indx, last_file_indx + 1):
+            new_datum_uri = self.get_datum_uri(indx)
+            new_asset = Asset(
+                data_uri=new_datum_uri,
+                is_directory=False,
+                parameter="data_uris",
+                num=len(self.assets) + 1,
+            )
+            self.assets.append(new_asset)
+            self.data_uris.append(new_datum_uri)
 
         super().consume_stream_datum(doc)
 


Test patch:
diff --git a/src/bluesky/tests/test_consolidators.py b/src/bluesky/tests/test_consolidators.py
index 511345ad..a78d1fab 100644
--- a/src/bluesky/tests/test_consolidators.py
+++ b/src/bluesky/tests/test_consolidators.py
@@ -1,6 +1,8 @@
+from math import ceil
+
 import pytest
 
-from bluesky.consolidators import HDF5Consolidator
+from bluesky.consolidators import HDF5Consolidator, TIFFConsolidator
 
 
 @pytest.fixture
@@ -41,7 +43,7 @@ def descriptor():
 
 
 @pytest.fixture
-def stream_resource_factory():
+def hdf5_stream_resource_factory():
     return lambda data_key, chunk_shape: {
         "data_key": data_key,
         "mimetype": "application/x-hdf5",
@@ -56,14 +58,25 @@ def stream_resource_factory():
     }
 
 
+@pytest.fixture
+def tiff_stream_resource_factory():
+    return lambda data_key, chunk_shape: {
+        "data_key": data_key,
+        "mimetype": "multipart/related;type=image/tiff",
+        "uri": "file://localhost/test/file/path",
+        "parameters": {"chunk_shape": chunk_shape, "template": "img_{:06d}.tiff"},
+        "uid": f"stream-resource-uid-{data_key}",
+    }
+
+
 @pytest.fixture
 def stream_datum_factory():
-    return lambda data_key, i: {
-        "seq_nums": {"start": i + 1, "stop": i + 2},
-        "indices": {"start": i, "stop": i + 1},
+    return lambda data_key, indx, i_start, i_stop: {
+        "seq_nums": {"start": i_start + 1, "stop": i_stop + 1},
+        "indices": {"start": i_start, "stop": i_stop},
         "descriptor": "descriptor-uid",
         "stream_resource": f"stream-resource-uid-{data_key}",
-        "uid": f"stream-datum-uid-{data_key}/{i}",
+        "uid": f"stream-datum-uid-{data_key}/{indx}",
     }
 
 
@@ -76,14 +89,31 @@ shape_testdata = [
 
 
 @pytest.mark.parametrize("data_key, expected", shape_testdata)
-def test_shape(descriptor, stream_resource_factory, stream_datum_factory, data_key, expected):
-    stream_resource = stream_resource_factory(data_key=data_key, chunk_shape=())
+def test_hdf5_shape(descriptor, hdf5_stream_resource_factory, stream_datum_factory, data_key, expected):
+    stream_resource = hdf5_stream_resource_factory(data_key=data_key, chunk_shape=())
     cons = HDF5Consolidator(stream_resource, descriptor)
     assert cons.shape == (0, *expected[1:])
     for i in range(5):
-        doc = stream_datum_factory(data_key, i)
+        doc = stream_datum_factory(data_key, i, i, i + 1)
+        cons.consume_stream_datum(doc)
+    assert cons.shape == expected
+
+
+@pytest.mark.parametrize("data_key, expected", shape_testdata)
+@pytest.mark.parametrize("files_per_stream_datum", [1, 2, 3, 5])
+def test_tiff_shape(
+    descriptor, tiff_stream_resource_factory, stream_datum_factory, data_key, expected, files_per_stream_datum
+):
+    stream_resource = tiff_stream_resource_factory(data_key=data_key, chunk_shape=(1,))
+    cons = TIFFConsolidator(stream_resource, descriptor)
+    assert cons.shape == (0, *expected[1:])
+    for i in range(ceil(5 / files_per_stream_datum)):
+        doc = stream_datum_factory(
+            data_key, i, i * files_per_stream_datum, min((i + 1) * files_per_stream_datum, 5)
+        )
         cons.consume_stream_datum(doc)
     assert cons.shape == expected
+    assert len(cons.assets) == 5
 
 
 chunk_testdata = [
@@ -102,11 +132,11 @@ chunk_testdata = [
 
 
 @pytest.mark.parametrize("data_key, chunk_shape, expected", chunk_testdata)
-def test_chunks(descriptor, stream_resource_factory, stream_datum_factory, data_key, chunk_shape, expected):
-    stream_resource = stream_resource_factory(data_key=data_key, chunk_shape=chunk_shape)
+def test_chunks(descriptor, hdf5_stream_resource_factory, stream_datum_factory, data_key, chunk_shape, expected):
+    stream_resource = hdf5_stream_resource_factory(data_key=data_key, chunk_shape=chunk_shape)
     cons = HDF5Consolidator(stream_resource, descriptor)
     assert cons.chunks == ((0,), *expected[1:])
     for i in range(5):
-        doc = stream_datum_factory(data_key, i)
+        doc = stream_datum_factory(data_key, i, i, i + 1)
         cons.consume_stream_datum(doc)
     assert cons.chunks == expected