Merge pull request #582 from European-XFEL/feat/more-magic-methods

Support more magic methods

extra_data/reader.py

@@ -290,6 +290,21 @@ def __getitem__(self, item):
 
         raise TypeError("Expected data[source], data[source, key] or data[train_selection]")
 
+    def __contains__(self, item):
+        if (
+            isinstance(item, tuple) and
+            len(item) == 2 and
+            all(isinstance(e, str) for e in item)
+        ):
+            return item[0] in self.all_sources and \
+                item[1] in self._get_source_data(item[0])
+        elif isinstance(item, str):
+            return item in self.all_sources
+
+        return False
+
+    __iter__ = None  # Disable iteration
+
     def _ipython_key_completions_(self):
         return list(self.all_sources)
 
@@ -650,6 +665,12 @@ def union(self, *others):
             is_single_run=same_run(self, *others),
         )
 
+    def __or__(self, other):
+        return self.union(other)
+
+    def __ior__(self, other):
+        return self.union(other)
+
     def _parse_aliases(self, alias_defs):
         """Parse alias definitions into alias dictionaries."""
 

extra_data/sourcedata.py

@@ -66,6 +66,8 @@ def __contains__(self, key):
             res = self._has_exact_key(key + '.value')
         return res
 
+    __iter__ = None  # Disable iteration
+
     def __getitem__(self, key):
         if (
             isinstance(key, (by_id, by_index, list, np.ndarray, slice)) or
@@ -501,3 +503,9 @@ def union(self, *others) -> 'SourceData':
             is_single_run=same_run(self, *others),
             inc_suspect_trains=self.inc_suspect_trains
         )
+
+    def __or__(self, other):
+        return self.union(other)
+
+    def __ior__(self, other):
+        return self.union(other)

extra_data/tests/test_reader_mockdata.py

@@ -843,6 +843,26 @@ def test_union(mock_fxe_raw_run):
     assert joined[xgm].train_ids == expected_tids
     assert joined[camera].train_ids == expected_tids
 
+    # Try via operators.
+    sel1 = run.select(xgm, 'beamPosition.ixPos')
+    sel2 = run.select(xgm, 'beamPosition.iyPos')
+
+    joined = sel1 | sel2
+    assert joined.selection == {
+        xgm: {
+            'beamPosition.ixPos.value',
+            'beamPosition.iyPos.value',
+        }
+    }
+
+    sel1 |= sel2
+    assert sel1.selection == {
+        xgm: {
+            'beamPosition.ixPos.value',
+            'beamPosition.iyPos.value',
+        }
+    }
+
 
 def test_union_raw_proc(mock_spb_raw_run, mock_spb_proc_run):
     raw_run = RunDirectory(mock_spb_raw_run)
@@ -1050,3 +1070,14 @@ def test_proc_legacy_sources(mock_modern_spb_proc_run):
         assert run.get_dtype(det_mod0, 'image.data') == np.float32
 
     assert run.select(det_mod0).all_sources == {det_mod0}
+
+
+def test_datacollection_contains(mock_fxe_control_data):
+    run = H5File(mock_fxe_control_data)
+
+    assert 'FXE_XAD_GEC/CAM/CAMERA:daqOutput' in run
+    assert 'MY/LITTLE/PONY' not in run
+    assert ('MY/LITTLE/PONY', 'actualPosition') not in run
+    assert ('SPB_XTD9_XGM/DOOCS/MAIN', 'beamPosition.ixPos') in run
+    assert ('SPB_XTD9/XGM/DOOCS/MAIN', '42') not in run
+    assert 42 not in run

extra_data/tests/test_sourcedata.py

@@ -155,6 +155,13 @@ def test_union(mock_spb_raw_run):
     with pytest.raises(ValueError):
         xgm.union(am0)
 
+    sel = xgm.select_trains(np.s_[:10]) | xgm.select_trains(np.s_[-10:])
+    assert sel.train_ids == list(range(10000, 10010)) + list(range(10054, 10064))
+
+    sel = xgm.select_trains(np.s_[:10])
+    sel |= xgm.select_trains(np.s_[-10:])
+    assert sel.train_ids == list(range(10000, 10010)) + list(range(10054, 10064))
+
 
 def test_run_value(mock_spb_raw_run):
     run = RunDirectory(mock_spb_raw_run)