Expose units symbol & name on KeyData objects

extra_data/keydata.py

@@ -106,6 +106,28 @@ def size_gb(self):
         """The size of the data in memory in gigabytes."""
         return self.nbytes / 1e9
 
+    @property
+    def units(self):
+        dset = self.files[0].file[self.hdf5_data_path]
+        base_unit = dset.attrs.get('unitSymbol', None)
+        if base_unit is None:
+            return None
+
+        prefix = dset.attrs.get('metricPrefixSymbol', '')
+        if prefix == 'u':
+            prefix = 'μ'  # We are not afraid of unicode
+        return prefix + base_unit
+
+    @property
+    def units_name(self):
+        dset = self.files[0].file[self.hdf5_data_path]
+        base_unit = dset.attrs.get('unitName', None)
+        if base_unit is None:
+            return None
+
+        prefix = dset.attrs.get('metricPrefixName', '')
+        return prefix + base_unit
+
     @property
     def source_file_paths(self):
         paths = []
@@ -144,8 +166,11 @@ def __getitem__(self, item):
 
     def _only_tids(self, tids):
         tids_arr = np.array(tids)
-        files = [f for f in self.files
-                 if f.has_train_ids(tids_arr, self.inc_suspect_trains)]
+        # Keep 1 file, even if 0 trains selected.
+        files = [
+            f for f in self.files
+            if f.has_train_ids(tids_arr, self.inc_suspect_trains)
+        ] or [self.files[0]]
 
         return KeyData(
             self.source,

extra_data/tests/mockdata/xgm.py

@@ -1,3 +1,5 @@
+import numpy as np
+
 from .base import DeviceBase
 
 class XGM(DeviceBase):
@@ -53,3 +55,16 @@ class XGM(DeviceBase):
         ('xTD', 'f4', (1000,)),
         ('yTD', 'f4', (1000,)),
     ]
+
+    def write_instrument(self, f):
+        super().write_instrument(f)
+
+        # Annotate intensityTD with some units to test retrieving them
+        # Karabo stores ASCII strings, assigning bytes is a shortcut to mimic that
+        ds = f[f'INSTRUMENT/{self.device_id}:output/data/intensityTD']
+        ds.attrs['metricPrefixEnum']= np.array([14], dtype=np.int32)
+        ds.attrs['metricPrefixName'] = b'micro'
+        ds.attrs['metricPrefixSymbol'] = b'u'
+        ds.attrs['unitEnum'] = np.array([15], dtype=np.int32)
+        ds.attrs['unitName'] = b'joule'
+        ds.attrs['unitSymbol'] = b'J'

extra_data/tests/test_keydata.py

@@ -49,7 +49,7 @@ def test_select_trains(mock_spb_raw_run):
     # Empty selection
     sel2 = xgm_beam_x[80:]
     assert sel2.shape == (0,)
-    assert len(sel2.files) == 0
+    assert len(sel2.files) == 1
     assert sel2.xarray().shape == (0,)
 
     # Single train
@@ -339,3 +339,14 @@ def test_file_no_trains(run_with_file_no_trains):
     run = RunDirectory(run_with_file_no_trains)
     xpos = run['SPB_XTD9_XGM/DOOCS/MAIN', 'beamPosition.ixPos'].ndarray()
     assert xpos.shape == (64,)
+
+
+def test_units(mock_sa3_control_data):
+    run = H5File(mock_sa3_control_data)
+    xgm_intensity = run['SA3_XTD10_XGM/XGM/DOOCS:output', 'data.intensityTD']
+
+    assert xgm_intensity.units == 'μJ'
+    assert xgm_intensity.units_name == 'microjoule'
+
+    # Check that it still works after selecting 0 trains
+    assert xgm_intensity.select_trains(np.s_[:0]).units == 'μJ'