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))
```

Discussion:

Patch:
diff --git a/.github/dependabot/constraints.txt b/.github/dependabot/constraints.txt
index ff0884c..8ca9cb1 100644
--- a/.github/dependabot/constraints.txt
+++ b/.github/dependabot/constraints.txt
@@ -13,7 +13,7 @@ psutil==5.8.0
 pyparsing==2.4.7
 python-dateutil==2.8.1
 pytz==2020.5
-pyzmq==21.0.1
-scipy<=1.5.4
+pyzmq==20.0.0
+scipy==1.5.4
 six==1.15.0
 xarray==0.16.2
diff --git a/docs/index.rst b/docs/index.rst
index b3f16e1..ae06c75 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -45,7 +45,7 @@ run or a file.
 Load data into memory - see :ref:`data-by-source-and-key` for more::
 
     # Get a labelled array
-    arr = run["SA3_XTD10_PES/ADC/1:network", "digitizers.channel_4_A.raw.samples"].xarray()
+    arr = run.get_array("SA3_XTD10_PES/ADC/1:network", "digitizers.channel_4_A.raw.samples")
 
     # Get a pandas dataframe of 1D fields
     df = run.get_dataframe(fields=[
diff --git a/docs/reading_files.rst b/docs/reading_files.rst
index 9e2abea..73664ec 100644
--- a/docs/reading_files.rst
+++ b/docs/reading_files.rst
@@ -71,71 +71,30 @@ to refer to all data associated with that 0.1 second window.
 Getting data by source & key
 ----------------------------
 
-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.::
+Where data will fit into memory, it's usually quickest and most convenient
+to load it like this.
 
-    xgm_intensity = run['SA1_XTD2_XGM/XGM/DOOCS:output', 'data.intensityTD'].xarray()
-
-.. class:: KeyData
-
-   .. attribute:: dtype
-
-      The NumPy dtype for this data. This indicates whether it contains
-      integers or floating point numbers, and how many bytes of memory each
-      number needs.
-
-   .. attribute:: ndim
-
-      The number of dimensions the data has. All data has at least 1 dimension
-      (time). A sequence of 2D images would have 3 dimensions.
-
-   .. autoattribute:: shape
-
-   .. attribute:: entry_shape
-
-      The shape of a single entry in the data, e.g. a single frame from a
-      camera. This is equivalent to ``key.shape[1:]``, but may be quicker than
-      that.
-
-   .. automethod:: data_counts
-
-   .. automethod:: ndarray
-
-   .. automethod:: series
+.. class:: DataCollection
+   :noindex:
 
-   .. automethod:: xarray
+   .. automethod:: get_array
 
      .. seealso::
        `xarray documentation <http://xarray.pydata.org/en/stable/indexing.html>`__
-         How to use the arrays returned by :meth:`DataCollection.get_array`
+         How to use the arrays returned by :meth:`~.get_array`
 
        :doc:`xpd_examples`
          Examples using xarray & pandas with EuXFEL data
 
-   .. automethod:: dask_array
+   .. automethod:: get_dask_array
 
-    .. seealso::
+     .. seealso::
        `Dask Array documentation <https://docs.dask.org/en/latest/array.html>`__
-         How to use the objects returned by :meth:`DataCollection.get_dask_array`
+         How to use the objects returned by :meth:`~.get_dask_array`
 
        :doc:`dask_averaging`
          An example using Dask with EuXFEL data
 
-   .. automethod:: select_trains
-
-The run or file object (a :class:`DataCollection`) also has methods to load
-data by sources and keys. Some of them are directly equivalent to the options
-above, but :meth:`.get_dataframe` and :meth:`.get_virtual_dataset` offer extra
-capabilities.
-
-.. class:: DataCollection
-   :noindex:
-
-   .. automethod:: get_array
-
-   .. automethod:: get_dask_array
-
    .. automethod:: get_series
 
    .. automethod:: get_dataframe
@@ -163,28 +122,9 @@ Getting data by train
 Some kinds of data, e.g. from AGIPD, are too big to load a whole run into
 memory at once. In these cases, it's convenient to load one train at a time.
 
-If you want to do this for just one source & key with :class:`KeyData` methods,
-like this::
-
-    for tid, arr in run['SA1_XTD2_XGM/XGM/DOOCS:output', 'data.intensityTD'].trains():
-        ...
-
-.. class:: KeyData
-   :noindex:
-
-   .. automethod:: trains
-
-   .. automethod:: train_from_id
-
-   .. automethod:: train_from_index
-
-To work with multiple modules of the same detector, see :doc:`agipd_lpd_data`.
-
-You can also get data by train for multiple sources and keys together from a run
-or file object.
-It's always a good idea to select the data you're interested in, either using
-:meth:`~.DataCollection.select`, or the ``devices=`` parameter. If you don't,
-they will read data for all sources in the run, which may be very slow.
+When accessing data like this, it's worth selecting which sources you're
+interested in, either using :meth:`~.DataCollection.select`, or the ``devices=``
+parameter. This avoids reading all the other data.
 
 .. class:: DataCollection
    :noindex:
@@ -280,7 +220,7 @@ Here are some problems we've seen, and possible solutions or workarounds:
 
   - In one case, a train ID had the maximum possible value (2\ :sup:`64` - 1),
     causing :meth:`~.info` to fail. You can select everything except this train
-    using :meth:`~.DataCollection.select_trains`::
+    using :meth:`~.select_trains`::
 
         from extra_data import by_id
         sel = run.select_trains(by_id[:2**64-1])
diff --git a/extra_data/__init__.py b/extra_data/__init__.py
index dfd12fb..71e8478 100644
--- a/extra_data/__init__.py
+++ b/extra_data/__init__.py
@@ -36,7 +36,6 @@ program. If not, see <https://opensource.org/licenses/BSD-3-Clause>
 __version__ = "1.3.0"
 
 
-from .keydata import KeyData
 from .reader import *
 from .stacking import *
 from .utils import *
diff --git a/extra_data/keydata.py b/extra_data/keydata.py
index 5a3c808..5e5573e 100644
--- a/extra_data/keydata.py
+++ b/extra_data/keydata.py
@@ -69,14 +69,10 @@ class KeyData:
 
     @property
     def shape(self):
-        """The shape of this data as a tuple, like for a NumPy array.
-
-        Finding the shape may require getting index data from several files
-        """
         return (sum(c.total_count for c in self._data_chunks),) + self.entry_shape
 
     def select_trains(self, trains):
-        """Select a subset of trains in this data as a new :class:`KeyData` object.
+        """Select a subset of the trains in this data
 
         Also available by slicing and indexing the KeyData object::
 
@@ -100,7 +96,7 @@ class KeyData:
         return self.select_trains(item)
 
     def data_counts(self):
-        """Get a count of data entries in each train.
+        """Get a count of data points in each train.
 
         Returns a pandas series with an index of train IDs.
         """
@@ -164,22 +160,11 @@ class KeyData:
         The first dimension is labelled with train IDs. Other dimensions may be
         named by passing a list of names to *extra_dims*.
 
-        Parameters
-        ----------
-
-        extra_dims: list of str
-            Name extra dimensions in the array. The first dimension is
-            automatically called 'train'. The default for extra dimensions
-            is dim_0, dim_1, ...
-        roi: numpy.s_[], slice, tuple of slices, or by_index
-            The region of interest. This expression selects data in all
-            dimensions apart from the first (trains) dimension. If the data
-            holds a 1D array for each entry, roi=np.s_[:8] would get the
-            first 8 values from every train. If the data is 2D or more at
-            each entry, selection looks like roi=np.s_[:8, 5:10] .
-        name: str
-            Name the array itself. The default is the source and key joined
-            by a dot.
+        *roi* may be a ``numpy.s_[]`` expression to load e.g. only part of each
+        image from a camera.
+
+        The array's name may be given by *name*, by default the source and key
+        is joined by a dot.
         """
         import xarray
 
@@ -204,8 +189,7 @@ class KeyData:
         return xarray.DataArray(ndarr, dims=dims, coords=coords, name=name)
 
     def series(self):
-        """Load this data as a pandas Series. Only for 1D data.
-        """
+        """Load this data as a pandas Series. Only for 1D data."""
         import pandas as pd
 
         if self.ndim > 1:
@@ -222,23 +206,11 @@ class KeyData:
     def dask_array(self, labelled=False):
         """Make a Dask array for this data.
 
-        Dask is a system for lazy parallel computation. This method doesn't
-        actually load the data, but gives you an array-like object which you
-        can operate on. Dask loads the data and calculates results when you ask
-        it to, e.g. by calling a ``.compute()`` method.
-        See the Dask documentation for more details.
-
-        If your computation depends on reading lots of data, consider creating
-        a dask.distributed.Client before calling this.
-        If you don't do this, Dask uses threads by default, which is not
-        efficient for reading HDF5 files.
-
-        Parameters
-        ----------
+        Dask does lazy computation, so the data is not actually loaded until
+        you call a ``.compute()`` method to get a result.
 
-        labelled: bool
-            If True, label the train IDs for the data, returning an
-            xarray.DataArray object wrapping a Dask array.
+        Passing ``labelled=True`` will wrap the Dask array in an xarray
+        DataArray with train IDs labelled.
         """
         import dask.array as da
 
diff --git a/extra_data/reader.py b/extra_data/reader.py
index 263c848..26957dc 100644
--- a/extra_data/reader.py
+++ b/extra_data/reader.py
@@ -527,9 +527,37 @@ class DataCollection:
         return pd.concat(series, axis=1)
 
     def get_array(self, source, key, extra_dims=None, roi=(), name=None):
-        """Return a labelled array for a data field defined by source and key.
+        """Return a labelled array for a particular data field.
 
-        see :meth:`.KeyData.xarray` for details.
+        ::
+
+            arr = run.get_array("SA3_XTD10_PES/ADC/1:network", "digitizers.channel_4_A.raw.samples")
+
+        This should work for any data.
+        The first axis of the returned data will be labelled with the train IDs.
+
+        Parameters
+        ----------
+
+        source: str
+            Device name with optional output channel, e.g.
+            "SA1_XTD2_XGM/DOOCS/MAIN" or "SPB_DET_AGIPD1M-1/DET/7CH0:xtdf"
+        key: str
+            Key of parameter within that device, e.g. "beamPosition.iyPos.value"
+            or "header.linkId".
+        extra_dims: list of str
+            Name extra dimensions in the array. The first dimension is
+            automatically called 'train'. The default for extra dimensions
+            is dim_0, dim_1, ...
+        roi: slice, tuple of slices, or by_index
+            The region of interest. This expression selects data in all
+            dimensions apart from the first (trains) dimension. If the data
+            holds a 1D array for each entry, roi=np.s_[:8] would get the
+            first 8 values from every train. If the data is 2D or more at
+            each entry, selection looks like roi=np.s_[:8, 5:10] .
+        name: str
+            Name the array itself. The default is the source and key joined
+            by a dot.
         """
         if isinstance(roi, by_index):
             roi = roi.value
@@ -538,9 +566,28 @@ class DataCollection:
             extra_dims=extra_dims, roi=roi, name=name)
 
     def get_dask_array(self, source, key, labelled=False):
-        """Get a Dask array for a data field defined by source and key.
+        """Get a Dask array for the specified data field.
+
+        Dask is a system for lazy parallel computation. This method doesn't
+        actually load the data, but gives you an array-like object which you
+        can operate on. Dask loads the data and calculates results when you ask
+        it to, e.g. by calling a ``.compute()`` method.
+        See the Dask documentation for more details.
+
+        If your computation depends on reading lots of data, consider creating
+        a dask.distributed.Client before calling this.
+        If you don't do this, Dask uses threads by default, which is not
+        efficient for reading HDF5 files.
 
-        see :meth:`.KeyData.dask_array` for details.
+        Parameters
+        ----------
+        source: str
+            Source name, e.g. "SPB_DET_AGIPD1M-1/DET/7CH0:xtdf"
+        key: str
+            Key of parameter within that device, e.g. "image.data".
+        labelled: bool
+            If True, label the train IDs for the data, returning an
+            xarray.DataArray object wrapping a Dask array.
         """
         return self._get_key_data(source, key).dask_array(labelled=labelled)
 
@@ -570,16 +617,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 +697,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',