Repository: European-XFEL/EXtra-data

Issue:
Selected keys can be overridden
`.select()` is meant to always make a narrower selection from the current one, but @daviddoji found a case where it can get broader again:

```python
source = 'FXE_DET_LPD1M-1/DET/2CH0:xtdf'

sel = run.select(source, "image.data")
print(sel.keys_for_source(source))
sel = sel.select(source, '*')
print(sel.keys_for_source(source))
```

Comments:

Patch:
diff --git a/extra_data/reader.py b/extra_data/reader.py
index 263c848..c388030 100644
--- a/extra_data/reader.py
+++ b/extra_data/reader.py
@@ -570,16 +570,54 @@ class DataCollection:
         res = defaultdict(set)
         if isinstance(selection, dict):
             # {source: {key1, key2}}
-            # {source: {}} -> all keys for this source
-            for source, keys in selection.items():  #
+            # {source: {}} or {source: None} -> all keys for this source
+            for source, in_keys in selection.items():
                 if source not in self.all_sources:
                     raise SourceNameError(source)
 
-                if keys:
-                    res[source].update(keys)
+                # Keys of the current DataCollection.
+                cur_keys = self.selection[source]
+
+                # Keys input as the new selection.
+                if in_keys:
+                    # 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):
+                            raise PropertyNameError(key, source)
                 else:
+                    # Catches both an empty set and None.
+                    # While the public API describes an empty set to
+                    # refer to all keys, the internal API actually uses
+                    # None for this case. This method is supposed to
+                    # accept both cases in order to natively support
+                    # passing a DataCollection as the selector. To keep
+                    # the conditions below clearer, any non-True value
+                    # is converted to None.
+                    in_keys = None
+
+                if cur_keys is None and in_keys is None:
+                    # Both the new and current keys select all.
                     res[source] = None
 
+                elif cur_keys is not None and in_keys is not None:
+                    # Both the new and current keys are specific, take
+                    # the intersection of both. This should never be
+                    # able to result in an empty set, but to prevent the
+                    # code further below from breaking, assert it.
+                    res[source] = cur_keys & in_keys
+                    assert res[source]
+
+                elif cur_keys is None and in_keys is not None:
+                    # Current keys are unspecific but new ones are, just
+                    # use the new keys.
+                    res[source] = in_keys
+
+                elif cur_keys is not None and in_keys is None:
+                    # The current keys are specific but new ones are
+                    # not, use the current keys.
+                    res[source] = cur_keys
+
         elif isinstance(selection, Iterable):
             # selection = [('src_glob', 'key_glob'), ...]
             res = union_selections(
@@ -612,7 +650,14 @@ class DataCollection:
                 continue
 
             if key_glob == '*':
-                matched[source] = None
+                # When the selection refers to all keys, make sure this
+                # is restricted to the current selection of keys for
+                # this source.
+
+                if self.selection[source] is None:
+                    matched[source] = None
+                else:
+                    matched[source] = self.selection[source]
             else:
                 r = ctrl_key_re if source in self.control_sources else key_re
                 keys = set(filter(r.match, self.keys_for_source(source)))


Test patch:
diff --git a/extra_data/tests/test_reader_mockdata.py b/extra_data/tests/test_reader_mockdata.py
index aa84f1b..8668764 100644
--- a/extra_data/tests/test_reader_mockdata.py
+++ b/extra_data/tests/test_reader_mockdata.py
@@ -476,17 +476,46 @@ def test_select(mock_fxe_raw_run):
 
     assert 'SPB_XTD9_XGM/DOOCS/MAIN' in run.control_sources
 
+    # Basic selection machinery, glob API
     sel = run.select('*/DET/*', 'image.pulseId')
     assert 'SPB_XTD9_XGM/DOOCS/MAIN' not in sel.control_sources
     assert 'FXE_DET_LPD1M-1/DET/0CH0:xtdf' in sel.instrument_sources
     _, data = sel.train_from_id(10000)
     for source, source_data in data.items():
-        print(source)
         assert set(source_data.keys()) == {'image.pulseId', 'metadata'}
 
-    sel2 = run.select(sel)
-    assert 'SPB_XTD9_XGM/DOOCS/MAIN' not in sel2.control_sources
-    assert 'FXE_DET_LPD1M-1/DET/0CH0:xtdf' in sel2.instrument_sources
+    # Basic selection machinery, dict-based API
+    sel_by_dict = run.select({
+        'SA1_XTD2_XGM/DOOCS/MAIN': None,
+        'FXE_DET_LPD1M-1/DET/0CH0:xtdf': {'image.pulseId'}
+    })
+    assert sel_by_dict.control_sources == {'SA1_XTD2_XGM/DOOCS/MAIN'}
+    assert sel_by_dict.instrument_sources == {'FXE_DET_LPD1M-1/DET/0CH0:xtdf'}
+    assert sel_by_dict.keys_for_source('FXE_DET_LPD1M-1/DET/0CH0:xtdf') == \
+        sel.keys_for_source('FXE_DET_LPD1M-1/DET/0CH0:xtdf')
+
+    # Re-select using * selection, should yield the same keys.
+    assert sel.keys_for_source('FXE_DET_LPD1M-1/DET/0CH0:xtdf') == \
+        sel.select('FXE_DET_LPD1M-1/DET/0CH0:xtdf', '*') \
+           .keys_for_source('FXE_DET_LPD1M-1/DET/0CH0:xtdf')
+    assert sel.keys_for_source('FXE_DET_LPD1M-1/DET/0CH0:xtdf') == \
+        sel.select({'FXE_DET_LPD1M-1/DET/0CH0:xtdf': {}}) \
+           .keys_for_source('FXE_DET_LPD1M-1/DET/0CH0:xtdf')
+
+    # Re-select a different but originally valid key, should fail.
+    with pytest.raises(ValueError):
+        # ValueError due to globbing.
+        sel.select('FXE_DET_LPD1M-1/DET/0CH0:xtdf', 'image.trainId')
+
+    with pytest.raises(PropertyNameError):
+        # PropertyNameError via explicit key.
+        sel.select({'FXE_DET_LPD1M-1/DET/0CH0:xtdf': {'image.trainId'}})
+
+    # Select by another DataCollection.
+    sel_by_dc = run.select(sel)
+    assert sel_by_dc.control_sources == sel.control_sources
+    assert sel_by_dc.instrument_sources == sel.instrument_sources
+    assert sel_by_dc.train_ids == sel.train_ids
 
 
 @pytest.mark.parametrize('select_str',