Merge pull request #85 from legend-exp/dev

Bug fix: support LH5 read/write of `array<n>{...}` where `n` > 1
legend-exp · May 1, 2024 · f26be5b · f26be5b
2 parents 04897fd + 7910df3
commit f26be5b
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 2 deletions.
diff --git a/src/lgdo/lh5/datatype.py b/src/lgdo/lh5/datatype.py
@@ -16,9 +16,9 @@
         ),
         (lgdo.Struct, r"^struct\{.*\}$"),
         (lgdo.Table, r"^table\{.*\}$"),
-        (lgdo.FixedSizeArray, r"^fixedsize_array<1>\{.+\}$"),
+        (lgdo.FixedSizeArray, r"^fixedsize_array<\d+>\{.+\}$"),
         (lgdo.ArrayOfEqualSizedArrays, r"^array_of_equalsized_arrays<1,1>\{.+\}$"),
-        (lgdo.Array, r"^array<1>\{.+\}$"),
+        (lgdo.Array, r"^array<\d+>\{.+\}$"),
     ]
 )
 """Mapping between LGDO types and regular expression defining the corresponding datatype string"""

diff --git a/tests/lh5/conftest.py b/tests/lh5/conftest.py
@@ -19,6 +19,7 @@ def lh5_file(tmptestdir):
     struct = lgdo.Struct()
     struct.add_field("scalar", lgdo.Scalar(value=10, attrs={"sth": 1}))
     struct.add_field("array", types.Array(nda=np.array([1, 2, 3, 4, 5])))
+    struct.add_field("array2d", types.Array(shape=(23, 56), fill_val=69, dtype=int))
     struct.add_field(
         "aoesa",
         types.ArrayOfEqualSizedArrays(shape=(5, 5), dtype=np.float32, fill_val=42),

diff --git a/tests/lh5/test_lh5_store.py b/tests/lh5/test_lh5_store.py
@@ -84,6 +84,10 @@ def test_read_array(lh5_file):
             is DEFAULT_HDF5_SETTINGS["compression"]
         )
 
+    lh5_obj, n_rows = store.read("/data/struct_full/array2d", lh5_file)
+    assert isinstance(lh5_obj, types.Array)
+    assert lh5_obj == types.Array(shape=(23, 56), fill_val=69, dtype=int)
+
 
 def test_read_array_slice(lh5_file):
     store = lh5.LH5Store()
@@ -351,6 +355,7 @@ def test_read_with_field_mask(lh5_file):
     assert sorted(lh5_obj.keys()) == [
         "aoesa",
         "array",
+        "array2d",
         "empty_struct",
         "scalar",
         "table",