Repository: European-XFEL/EXtra-data

Issue:
Objects for individual sources or datasets?
The `DataCollection` object is the centrepiece of EXtra-data, and has a tendency to keep growing. Many of its methods take `source, key` parameters to refer to a specific dataset, e.g. `get_array()` or `get_data_counts()`. There are others possible - e.g. how do you find the dtype or entry shape for a given key? A couple take just `source`, e.g. `keys_for_source()`.

Should we add new types to represent an individual source or an individual dataset in the available data? That would allow APIs maybe like this:

```python
run.keys_for_source("SA3_XTD10_PES/ADC/1:network")  # Current
run["SA3_XTD10_PES/ADC/1:network"].keys()  # Hypothetical

# Current:
run.get_array("SA3_XTD10_PES/ADC/1:network", "digitizers.channel_4_A.raw.samples")
# Hypothetical:
run["SA3_XTD10_PES/ADC/1:network", "digitizers.channel_4_A.raw.samples"].get_array()
```

And it would give a convenient place to expose things like the dtype & entry shape of a key, without awkward methods on DataCollection.

My reservations are:

- Do more classes make it harder to understand? I recall that one of the original goals was to use basic data structures like dicts as far as possible.
- The naming challenge - I've referred to 'sources' and 'datasets' here, but e.g. 'dataset' already means one thing in HDF5 and a different thing in Xarray, so that's not a great name. 'Source' also has many meanings, although not so immediately adjacent.

@tmichela @philsmt & anyone else who's interested: thoughts? This is for discussion at the moment, I'm not asking anyone to start coding.

Discussion:

Patch:
diff --git a/docs/reading_files.rst b/docs/reading_files.rst
index a048f5a..7417947 100644
--- a/docs/reading_files.rst
+++ b/docs/reading_files.rst
@@ -78,6 +78,13 @@ to refer to all data associated with that 0.1 second window.
 Getting data by source & key
 ----------------------------
 
+Selecting a source in a run gives a :class:`.SourceData` object.
+You can use this to find keys belonging to that source::
+
+    xgm = run['SPB_XTD9_XGM/DOOCS/MAIN']
+    xgm.keys()  # List the available keys
+    beam_x = xgm['beamPosition.ixPos']  # Get a KeyData object
+
 Selecting a single source & key in a run gives a :class:`.KeyData` object.
 You can get the data from this in various forms with the methods described
 below, e.g.::
diff --git a/extra_data/__init__.py b/extra_data/__init__.py
index 34dedfe..fa73451 100644
--- a/extra_data/__init__.py
+++ b/extra_data/__init__.py
@@ -41,8 +41,11 @@ from .exceptions import (
 )
 from .keydata import KeyData
 from .reader import *
+from .sourcedata import SourceData
 from .stacking import *
 from .utils import *
 
 
-__all__ = reader.__all__ + utils.__all__ + stacking.__all__
+__all__ = reader.__all__ + utils.__all__ + stacking.__all__ + [
+    'SourceData', 'KeyData'
+]
diff --git a/extra_data/file_access.py b/extra_data/file_access.py
index 4b43440..5508d03 100644
--- a/extra_data/file_access.py
+++ b/extra_data/file_access.py
@@ -381,7 +381,7 @@ class FileAccess:
                 res.add(key.replace('/', '.'))
 
         self.file['/RUN/' + source].visititems(add_key)
-        self._keys_cache[source] = res
+        self._run_keys_cache[source] = res
         return res
 
     def has_source_key(self, source, key):
diff --git a/extra_data/keydata.py b/extra_data/keydata.py
index a58522d..86fc141 100644
--- a/extra_data/keydata.py
+++ b/extra_data/keydata.py
@@ -60,7 +60,7 @@ class KeyData:
 
     def __repr__(self):
         return f"<extra_data.KeyData source={self.source!r} key={self.key!r} " \
-               f"for {len(self.train_ids)} trains"
+               f"for {len(self.train_ids)} trains>"
 
     @property
     def _key_group(self):
diff --git a/extra_data/reader.py b/extra_data/reader.py
index a1f2501..a1806d2 100644
--- a/extra_data/reader.py
+++ b/extra_data/reader.py
@@ -46,6 +46,7 @@ from .read_machinery import (
     find_proposal,
 )
 from .run_files_map import RunFilesMap
+from .sourcedata import SourceData
 from . import locality
 from .file_access import FileAccess
 
@@ -206,17 +207,9 @@ class DataCollection:
         if source not in self.all_sources:
             raise SourceNameError(source)
 
-        if not self._has_source_key(source, key):
+        if key not in self[source]:
             raise PropertyNameError(key, source)
 
-    def _has_source_key(self, source, key):
-        selected_keys = self.selection[source]
-        if selected_keys is not None:
-            return key in selected_keys
-
-        for f in self._source_index[source]:
-            return f.has_source_key(source, key)
-
     def keys_for_source(self, source):
         """Get a set of key names for the given source
 
@@ -228,37 +221,37 @@ class DataCollection:
         combine two runs where the source was configured differently, the
         result can be unpredictable.
         """
-        selected_keys = self.selection[source]
-        if selected_keys is not None:
-            return selected_keys
-
-        # The same source may be in multiple files, but this assumes it has
-        # the same keys in all files that it appears in.
-        for f in self._source_index[source]:
-            return f.get_keys(source)
+        return self._get_source_data(source).keys()
 
     # Leave old name in case anything external was using it:
     _keys_for_source = keys_for_source
 
     def _get_key_data(self, source, key):
-        self._check_field(source, key)
-        section = 'INSTRUMENT' if source in self.instrument_sources else 'CONTROL'
+        return self._get_source_data(source)[key]
+
+    def _get_source_data(self, source):
+        if source in self.control_sources:
+            section = 'CONTROL'
+        elif source in self.instrument_sources:
+            section = 'INSTRUMENT'
+        else:
+            raise SourceNameError(source)
+
         files = self._source_index[source]
-        ds0 = files[0].file[f"{section}/{source}/{key.replace('.', '/')}"]
-        return KeyData(
+        return SourceData(
             source,
-            key,
+            sel_keys=self.selection[source],
             train_ids=self.train_ids,
             files=files,
             section=section,
-            dtype=ds0.dtype,
-            eshape=ds0.shape[1:],
             inc_suspect_trains=self.inc_suspect_trains,
         )
 
     def __getitem__(self, item):
         if isinstance(item, tuple) and len(item) == 2:
             return self._get_key_data(*item)
+        elif isinstance(item, str):
+            return self._get_source_data(item)
 
         raise TypeError("Expected data[source, key]")
 
@@ -623,7 +616,7 @@ class DataCollection:
                     # If a specific set of keys is selected, make sure
                     # they are all valid.
                     for key in in_keys:
-                        if not self._has_source_key(source, key):
+                        if key not in self[source]:
                             raise PropertyNameError(key, source)
                 else:
                     # Catches both an empty set and None.
@@ -1064,11 +1057,12 @@ class DataCollection:
             print('  -', s)
             if any(p.match(s) for p in details_sources_re):
                 # Detail for control sources: list keys
-                ctrl_keys = self.keys_for_source(s)
+                ctrl_keys = self[s].keys(inc_timestamps=False)
                 print('    - Control keys (1 entry per train):')
                 keys_detail(s, sorted(ctrl_keys), prefix='      - ')
 
                 run_keys = self._source_index[s][0].get_run_keys(s)
+                run_keys = {k[:-6] for k in run_keys if k.endswith('.value')}
                 run_only_keys = run_keys - ctrl_keys
                 if run_only_keys:
                     print('    - Additional run keys (1 entry per run):')
diff --git a/extra_data/sourcedata.py b/extra_data/sourcedata.py
new file mode 100644
index 0000000..3ef8e81
--- /dev/null
+++ b/extra_data/sourcedata.py
@@ -0,0 +1,89 @@
+from typing import List
+
+import h5py
+
+from .exceptions import PropertyNameError
+from .file_access import FileAccess
+from .keydata import KeyData
+
+class SourceData:
+    """Data for one key in one source
+
+    Don't create this directly; get it from ``run[source]``.
+    """
+    def __init__(
+            self, source, *, sel_keys, train_ids, files, section,
+            inc_suspect_trains=True,
+    ):
+        self.source = source
+        self.sel_keys = sel_keys
+        self.train_ids = train_ids
+        self.files: List[FileAccess] = files
+        self.section = section
+        self.inc_suspect_trains = inc_suspect_trains
+
+    def __repr__(self):
+        return f"<extra_data.SourceData source={self.source!r} " \
+               f"for {len(self.train_ids)} trains>"
+
+    def _has_exact_key(self, key):
+        if self.sel_keys is not None:
+            return key in self.sel_keys
+
+        for f in self.files:
+            return f.has_source_key(self.source, key)
+
+    def __contains__(self, key):
+        res = self._has_exact_key(key)
+        if (not res) and (self.section == 'CONTROL'):
+            res = self._has_exact_key(key + '.value')
+        return res
+
+    def __getitem__(self, key):
+        if key not in self:
+            raise PropertyNameError(key, self.source)
+        ds0 = self.files[0].file[
+            f"{self.section}/{self.source}/{key.replace('.', '/')}"
+        ]
+        if isinstance(ds0, h5py.Group):
+            # This can only occur with a CONTROL key missing its .value suffix
+            ds0 = ds0['value']
+            key += '.value'
+
+        return KeyData(
+            self.source,
+            key,
+            train_ids=self.train_ids,
+            files=self.files,
+            section=self.section,
+            dtype=ds0.dtype,
+            eshape=ds0.shape[1:],
+            inc_suspect_trains=self.inc_suspect_trains,
+        )
+
+    def keys(self, inc_timestamps=True):
+        """Get a set of key names for this source
+
+        If you have used :meth:`select` to filter keys, only selected keys
+        are returned.
+
+        For control sources, each Karabo property is stored in the file as two
+        keys, with '.value' and '.timestamp' suffixes. By default, these are
+        given separately. Pass ``inc_timestamps=False`` to ignore timestamps and
+        drop the '.value' suffix, giving names as used in Karabo.
+
+        Only one file is used to find the keys. Within a run, all files should
+        have the same keys for a given source, but if you use :meth:`union` to
+        combine two runs where the source was configured differently, the
+        result can be unpredictable.
+        """
+        if (not inc_timestamps) and (self.section == 'CONTROL'):
+            return {k[:-6] for k in self.keys() if k.endswith('.value')}
+
+        if self.sel_keys is not None:
+            return self.sel_keys
+
+        # The same source may be in multiple files, but this assumes it has
+        # the same keys in all files that it appears in.
+        for f in self.files:
+            return f.get_keys(self.source)


Test patch:
diff --git a/extra_data/tests/test_sourcedata.py b/extra_data/tests/test_sourcedata.py
new file mode 100644
index 0000000..3e92fd3
--- /dev/null
+++ b/extra_data/tests/test_sourcedata.py
@@ -0,0 +1,31 @@
+import numpy as np
+
+from extra_data import RunDirectory
+
+def test_get_sourcedata(mock_spb_raw_run):
+    run = RunDirectory(mock_spb_raw_run)
+    am0 = run['SPB_DET_AGIPD1M-1/DET/0CH0:xtdf']
+    assert len(am0.files) == 1
+    assert am0.section == 'INSTRUMENT'
+
+    xgm = run['SPB_XTD9_XGM/DOOCS/MAIN']
+    assert len(xgm.files) == 2
+    assert xgm.section == 'CONTROL'
+
+def test_keys(mock_spb_raw_run):
+    run = RunDirectory(mock_spb_raw_run)
+    xgm = run['SPB_XTD9_XGM/DOOCS/MAIN']
+
+    # Control keys can omit .value suffix, but .keys() will not list that.
+    assert 'beamPosition.ixPos.value' in xgm
+    assert 'beamPosition.ixPos' in xgm
+    assert 'beamPosition.ixPos.value' in xgm.keys()
+    assert 'beamPosition.ixPos.timestamp' in xgm.keys()
+    assert 'beamPosition.ixPos' not in xgm.keys()
+    assert xgm['beamPosition.ixPos.value'].dtype == np.dtype('f4')
+    assert xgm['beamPosition.ixPos'].dtype == np.dtype('f4')
+
+    # .keys(inc_timestamp=False) will give us only the name before '.value'
+    assert 'beamPosition.ixPos.value' not in xgm.keys(inc_timestamps=False)
+    assert 'beamPosition.ixPos.timestamp' not in xgm.keys(inc_timestamps=False)
+    assert 'beamPosition.ixPos' in xgm.keys(inc_timestamps=False)