Merge pull request #2451 from cta-observatory/path_input_h5_datalevels

Enable str/Path as input to get_hdf5_datalevels

Author: maxnoe

ctapipe/io/hdf5eventsource.py

@@ -1,7 +1,8 @@
 import logging
+from contextlib import ExitStack
 from functools import lru_cache
 from pathlib import Path
-from typing import Dict
+from typing import Dict, Union
 
 import astropy.units as u
 import numpy as np
@@ -76,24 +77,28 @@
 ]
 
 
-def get_hdf5_datalevels(h5file):
+def get_hdf5_datalevels(h5file: Union[tables.File, str, Path]):
     """Get the data levels present in the hdf5 file"""
     datalevels = []
 
-    if "/r1/event/telescope" in h5file.root:
-        datalevels.append(DataLevel.R1)
+    with ExitStack() as stack:
+        if not isinstance(h5file, tables.File):
+            h5file = stack.enter_context(tables.open_file(h5file))
 
-    if "/dl1/event/telescope/images" in h5file.root:
-        datalevels.append(DataLevel.DL1_IMAGES)
+        if "/r1/event/telescope" in h5file.root:
+            datalevels.append(DataLevel.R1)
 
-    if "/dl1/event/telescope/parameters" in h5file.root:
-        datalevels.append(DataLevel.DL1_PARAMETERS)
+        if "/dl1/event/telescope/images" in h5file.root:
+            datalevels.append(DataLevel.DL1_IMAGES)
 
-    if "/dl1/event/telescope/muon" in h5file.root:
-        datalevels.append(DataLevel.DL1_MUON)
+        if "/dl1/event/telescope/parameters" in h5file.root:
+            datalevels.append(DataLevel.DL1_PARAMETERS)
 
-    if "/dl2" in h5file.root:
-        datalevels.append(DataLevel.DL2)
+        if "/dl1/event/telescope/muon" in h5file.root:
+            datalevels.append(DataLevel.DL1_MUON)
+
+        if "/dl2" in h5file.root:
+            datalevels.append(DataLevel.DL2)
 
     return tuple(datalevels)
 

ctapipe/io/tests/test_hdf5.py

@@ -1,4 +1,5 @@
 import enum
+from pathlib import Path
 
 import numpy as np
 import pandas as pd
@@ -17,6 +18,7 @@
 )
 from ctapipe.core.container import Container, Field
 from ctapipe.io import read_table
+from ctapipe.io.datalevels import DataLevel
 from ctapipe.io.hdf5tableio import HDF5TableReader, HDF5TableWriter
 
 
@@ -1014,3 +1016,21 @@ class Container1(Container):
         assert c1.prefix == "bar"
         assert c2.value == value
         assert c2.prefix == "foo"
+
+
+@pytest.mark.parametrize("input_type", (str, Path, tables.File))
+def test_hdf5_datalevels(input_type, dl2_shower_geometry_file):
+    from ctapipe.io import get_hdf5_datalevels
+
+    if input_type is tables.File:
+        with tables.open_file(dl2_shower_geometry_file) as h5file:
+            datalevels = get_hdf5_datalevels(h5file)
+    else:
+        path = input_type(dl2_shower_geometry_file)
+        datalevels = get_hdf5_datalevels(path)
+
+    assert set(datalevels) == {
+        DataLevel.DL1_IMAGES,
+        DataLevel.DL1_PARAMETERS,
+        DataLevel.DL2,
+    }

docs/changes/2451.feature.rst

@@ -0,0 +1,2 @@
+Add support for using ``str`` and ``Path`` objects as input
+to ``ctapipe.io.get_hdf5_datalevels``.