Repository: European-XFEL/EXtra-data

Issue:
by_id and by_index slices do not support scalar integers
Unlike `np.s_`, our own special slices `by_id` and `by_index` do not support scalar integers to select a single train by ID or index:

* `run.select_trains(np.s_[0])`
    ```python
    <extra_data.DataCollection for 81 sources and 1 trains>
    ``` 

* `run.select_trains(by_index[0])`
    ```python
    ---------------------------------------------------------------------------
  TypeError                                 Traceback (most recent call last)
  Cell In[130], line 1
  ----> 1 run.select_trains(by_index[0])
  
  File ~/code/EXtra-data/extra_data/reader.py:1091, in DataCollection.select_trains(self, train_range)
     1071 def select_trains(self, train_range):
     1072     """Select a subset of trains from this data.
     1073 
     1074     Slice trains by position within this data::
     (...)
     1089         If given train IDs do not overlap with the trains in this data.
     1090     """
  -> 1091     new_train_ids = select_train_ids(self.train_ids, train_range)
     1093     sources_data = {
     1094         src: srcdata._only_tids(new_train_ids)
     1095         for src, srcdata in self._sources_data.items()
     1096     }
     1098     files = set().union(*[sd.files for sd in sources_data.values()])
  
  File ~/code/EXtra-data/extra_data/read_machinery.py:135, in select_train_ids(train_ids, sel)
      133     return sorted(np.asarray(train_ids)[sel])
      134 else:
  --> 135     raise TypeError(type(sel))
  
  TypeError: <class 'int'>
    ```


* `run.select_trains(by_id[run.train_ids[0]])`
    ```python
    ---------------------------------------------------------------------------
  TypeError                                 Traceback (most recent call last)
  Cell In[131], line 1
  ----> 1 run.select_trains(by_id[run.train_ids[0]])
  
  File ~/code/EXtra-data/extra_data/reader.py:1091, in DataCollection.select_trains(self, train_range)
     1071 def select_trains(self, train_range):
     1072     """Select a subset of trains from this data.
     1073 
     1074     Slice trains by position within this data::
     (...)
     1089         If given train IDs do not overlap with the trains in this data.
     1090     """
  -> 1091     new_train_ids = select_train_ids(self.train_ids, train_range)
     1093     sources_data = {
     1094         src: srcdata._only_tids(new_train_ids)
     1095         for src, srcdata in self._sources_data.items()
     1096     }
     1098     files = set().union(*[sd.files for sd in sources_data.values()])
  
  File ~/code/EXtra-data/extra_data/read_machinery.py:135, in select_train_ids(train_ids, sel)
      133     return sorted(np.asarray(train_ids)[sel])
      134 else:
  --> 135     raise TypeError(type(sel))
  
  TypeError: <class 'extra_data.read_machinery.by_id'>
    ```

Discussion:

Solution patch:
diff --git a/extra_data/read_machinery.py b/extra_data/read_machinery.py
index 835f8ea4..c3aefc36 100644
--- a/extra_data/read_machinery.py
+++ b/extra_data/read_machinery.py
@@ -4,11 +4,13 @@
 """
 import logging
 import math
+import operator
 import os
 import os.path as osp
 import re
 import time
 from glob import iglob
+from numbers import Integral
 from warnings import warn
 
 import numpy as np
@@ -106,29 +108,38 @@ def _tid_to_slice_ix(tid, train_ids, stop=False):
         return (train_ids > tid).nonzero()[0][0]
 
 
+def is_int_like(x):
+    if isinstance(x, np.ndarray):
+        return x.ndim == 0 and np.issubdtype(x.dtype, np.integer)
+    return isinstance(x, Integral)
+
+
 def select_train_ids(train_ids, sel):
     if isinstance(sel, by_index):
         sel = sel.value
-    elif isinstance(sel, int):
-        sel = slice(sel, sel+1, None)
-
-    if isinstance(sel, by_id) and isinstance(sel.value, slice):
-        # Slice by train IDs
-        start_ix = _tid_to_slice_ix(sel.value.start, train_ids, stop=False)
-        stop_ix = _tid_to_slice_ix(sel.value.stop, train_ids, stop=True)
-        return train_ids[start_ix: stop_ix: sel.value.step]
-    elif isinstance(sel, by_id) and isinstance(sel.value, (list, np.ndarray)):
-        # Select a list of trains by train ID
-        new_train_ids = sorted(set(train_ids).intersection(sel.value))
-        if len(sel.value) and not new_train_ids:
-            warn(
-                f"Given train IDs not found among {len(train_ids)} trains in "
-                "collection", stacklevel=3,
-            )
-        return new_train_ids
+
+    if isinstance(sel, by_id):
+        if isinstance(sel.value, slice):
+            # Slice by train IDs
+            start_ix = _tid_to_slice_ix(sel.value.start, train_ids, stop=False)
+            stop_ix = _tid_to_slice_ix(sel.value.stop, train_ids, stop=True)
+            return train_ids[start_ix: stop_ix: sel.value.step]
+        if is_int_like(sel.value):
+            sel.value = [operator.index(sel.value)]
+        if isinstance(sel.value, (list, np.ndarray)):
+            # Select a list of trains by train ID
+            new_train_ids = sorted(set(train_ids).intersection(sel.value))
+            if len(sel.value) and not new_train_ids:
+                warn(f"Given train IDs not found among {len(train_ids)} trains"
+                     " in collection", stacklevel=3,)
+            return new_train_ids
+        else:
+            raise TypeError(type(sel.value))
     elif isinstance(sel, slice):
         # Slice by indexes in this collection
         return train_ids[sel]
+    elif is_int_like(sel):
+        return [train_ids[operator.index(sel)]]
     elif isinstance(sel, (list, np.ndarray)):
         # Select a list of trains by index in this collection
         return sorted(np.asarray(train_ids)[sel])


Test patch:
diff --git a/extra_data/tests/test_read_machinery.py b/extra_data/tests/test_read_machinery.py
index 5f774a41..54271bb3 100644
--- a/extra_data/tests/test_read_machinery.py
+++ b/extra_data/tests/test_read_machinery.py
@@ -3,7 +3,10 @@
 from unittest import mock
 
 import numpy as np
-from extra_data import RunDirectory, read_machinery
+import pytest
+
+from extra_data import RunDirectory, by_id, by_index, read_machinery
+from extra_data.read_machinery import select_train_ids
 
 
 def test_find_proposal(tmpdir):
@@ -51,3 +54,77 @@ def test_same_run(mock_spb_raw_run, mock_jungfrau_run, mock_scs_run):
         assert read_machinery.same_run(sd1, sd2)
     else:
         assert not read_machinery.same_run(sd1, sd2)
+
+
+def test_select_train_ids():
+    train_ids = list(range(1000000, 1000010))
+
+    # Test by_id with a single integer
+    assert select_train_ids(train_ids, by_id[1000002]) == [1000002]
+
+    # Test by_id with a numpy type
+    assert select_train_ids(train_ids, by_id[np.uint64(1000002)]) == [1000002]
+
+    # Test by_id with a slice
+    assert select_train_ids(train_ids, by_id[1000002:1000005]) == [1000002, 1000003, 1000004]
+
+    # Test by_id with a list
+    assert select_train_ids(train_ids, by_id[[1000002, 1000005]]) == [1000002, 1000005]
+
+    # Test by_id with a numpy array
+    assert select_train_ids(train_ids, by_id[np.array([1000002, 1000005])]) == [1000002, 1000005]
+
+    # Test by_id with a slice and step
+    assert select_train_ids(train_ids, by_id[1000000:1000008:2]) == [1000000, 1000002, 1000004, 1000006]
+
+    # Test by_id with an open-ended slice (end)
+    assert select_train_ids(train_ids, by_id[1000005:]) == [1000005, 1000006, 1000007, 1000008, 1000009]
+
+    # Test by_id with an open-ended slice (start)
+    assert select_train_ids(train_ids, by_id[:1000003]) == [1000000, 1000001, 1000002]
+
+    # Test by_index with a single integer
+    assert select_train_ids(train_ids, by_index[2]) == [1000002]
+
+    # Test by_index with a 0D numpy array
+    assert select_train_ids(train_ids, by_index[np.array(5, dtype=np.int8)]) == [1000005]
+
+    # Test by_index with a slice
+    assert select_train_ids(train_ids, by_index[1:4]) == [1000001, 1000002, 1000003]
+
+    # Test by_index with a list
+    assert select_train_ids(train_ids, by_index[[1, 3]]) == [1000001, 1000003]
+
+    # Test by_index with a slice and step
+    assert select_train_ids(train_ids, by_index[::2]) == [1000000, 1000002, 1000004, 1000006, 1000008]
+
+    # Test with a plain integer
+    assert select_train_ids(train_ids, 3) == [1000003]
+
+    # Test with a plain int-like
+    assert select_train_ids(train_ids, np.uint32(3)) == [1000003]
+
+    # Test with a plain slice
+    assert select_train_ids(train_ids, slice(1, 4)) == [1000001, 1000002, 1000003]
+
+    # Test with a plain list
+    assert select_train_ids(train_ids, [1, 3]) == [1000001, 1000003]
+
+    # Test with a numpy array
+    assert select_train_ids(train_ids, np.array([1, 3])) == [1000001, 1000003]
+
+    # Test with an invalid type (should raise TypeError)
+    with pytest.raises(TypeError):
+        select_train_ids(train_ids, "invalid")
+    
+    with pytest.raises(TypeError):
+        select_train_ids(train_ids, by_id[np.float64(1000006)])
+
+    # Test by_id with train IDs not found in the list (should raise a warning)
+    with pytest.warns(UserWarning):
+        result = select_train_ids(train_ids, by_id[[999999, 1000010]])
+    assert result == []
+    
+    with pytest.warns(UserWarning):
+        result = select_train_ids(train_ids, by_id[1000010])
+    assert result == []