Add SourceData.run_keys()

Author: philsmt

extra_data/sourcedata.py

@@ -493,6 +493,29 @@ def run_value(self, key, *, allow_multi_run=False):
             return val.decode('utf-8', 'surrogateescape')
         return val
 
+    def run_keys(self, inc_timestamps=True):
+        """Get a set of RUN keys for this source.
+
+        Unlike :meth:`keys`, RUN keys are not affected by selection.
+        """
+
+        if not self.is_single_run:
+            raise MultiRunError()
+
+        if self.source not in self.files[0].file['RUN']:
+            raise ValueError(f'{self.source} has no RUN values')
+
+        res = set()
+        def visitor(path, obj):
+            if isinstance(obj, h5py.Dataset):
+                res.add(path.replace('/', '.'))
+
+        # Arbitrary file - should be the same across a run
+        self.files[0].file['RUN'][self.source].visititems(visitor)
+        if not inc_timestamps:
+            return {k[:-6] for k in res if k.endswith('.value')}
+        return res
+
     def run_values(self, inc_timestamps=True):
         """Get a dict of all RUN values for this source
 

extra_data/tests/test_sourcedata.py

@@ -186,6 +186,9 @@ def test_run_value(mock_spb_raw_run):
     assert 'pulseEnergy.conversion' in values_dict
     assert 'pulseEnergy.conversion.timestamp' not in values_dict
 
+    assert xgm.run_keys() == xgm.keys() | {'classId.timestamp', 'classId.value'}
+    assert xgm.run_keys(False) == xgm.keys(False) | {'classId'}
+
     with pytest.raises(ValueError):
         # no run values for AGIPD instrument data.
         am0.run_values()