Selecting train IDs in DataCollection and SourceData

extra_data/reader.py

@@ -33,17 +33,17 @@
 import numpy as np
 
 from . import locality, voview
+from .aliases import AliasIndexer
 from .exceptions import (MultiRunError, PropertyNameError, SourceNameError,
                          TrainIDError)
 from .file_access import FileAccess
 from .keydata import KeyData
-from .read_machinery import (DETECTOR_SOURCE_RE, FilenameInfo, by_id, by_index,
-                             find_proposal, glob_wildcards_re, same_run,
-                             select_train_ids, split_trains)
+from .read_machinery import (DETECTOR_SOURCE_RE, by_id, by_index,
+                             find_proposal, glob_wildcards_re, is_int_like,
+                             same_run, select_train_ids)
 from .run_files_map import RunFilesMap
 from .sourcedata import SourceData
 from .utils import available_cpu_cores
-from .aliases import AliasIndexer
 
 __all__ = [
     'H5File',
@@ -278,8 +278,13 @@ def __getitem__(self, item):
             return self._get_key_data(*item)
         elif isinstance(item, str):
             return self._get_source_data(item)
+        elif (
+            isinstance(item, (by_id, by_index, list, np.ndarray, slice)) or
+            is_int_like(item)
+        ):
+            return self.select_trains(item)
 
-        raise TypeError("Expected data[source] or data[source, key]")
+        raise TypeError("Expected data[source], data[source, key] or data[train_selection]")
 
     def _ipython_key_completions_(self):
         return list(self.all_sources)

extra_data/sourcedata.py

@@ -2,15 +2,15 @@
 import re
 from typing import Dict, List, Optional
 
-import numpy as np
 import h5py
+import numpy as np
 
-from .exceptions import MultiRunError, PropertyNameError, NoDataError
+from .exceptions import MultiRunError, NoDataError, PropertyNameError
 from .file_access import FileAccess
 from .keydata import KeyData
-from .read_machinery import (
-    glob_wildcards_re, same_run, select_train_ids, split_trains, trains_files_index
-)
+from .read_machinery import (by_id, by_index, glob_wildcards_re, is_int_like,
+                             same_run, select_train_ids, split_trains,
+                             trains_files_index)
 
 
 class SourceData:
@@ -67,6 +67,12 @@ def __contains__(self, key):
         return res
 
     def __getitem__(self, key):
+        if (
+            isinstance(key, (by_id, by_index, list, np.ndarray, slice)) or
+            is_int_like(key)
+        ):
+            return self.select_trains(key)
+
         if key not in self:
             raise PropertyNameError(key, self.source)
         ds0 = self.files[0].file[

extra_data/tests/test_reader_mockdata.py

@@ -717,6 +717,12 @@
     with pytest.raises(IndexError):
         run.select_trains(by_index[[480]])
 
+    assert run[10].train_ids == [10010]
+    assert run[by_id[10000]].train_ids == [10000]
+    assert run[by_index[479:555]].train_ids == [10479]
+    with pytest.raises(IndexError):
+        run[555]
+
 
 def test_split_trains(mock_fxe_raw_run):
     run = RunDirectory(mock_fxe_raw_run)

extra_data/tests/test_sourcedata.py

@@ -114,6 +114,19 @@ def test_select_trains(mock_spb_raw_run):
     assert sel.train_ids == []
     assert sel.keys() == xgm.keys()
 
+    sel = xgm[by_id[10020:10040]]
+    assert sel.train_ids == list(range(10020, 10040))
+
+    sel = xgm[by_index[:10]]
+    assert sel.train_ids == list(range(10000, 10010))
+
+    sel = xgm[10]
+    assert sel.train_ids == [10010]
+
+    sel = xgm[999:1000]
+    assert sel.train_ids == []
+    assert sel.keys() == xgm.keys()
+
 
 def test_split_trains(mock_spb_raw_run):
     run = RunDirectory(mock_spb_raw_run)