Use _select_subcomponent_trains() helper in Timepix3.select_trains()

Author: takluyver

src/extra/components/las.py

@@ -288,7 +288,7 @@ def select_trains(self, trains):
             '_stage_src', '_stage_delay_key',
             '_bam_src', '_bam_delay_key',
             '_pulses'
-        ])
+        ], trains)
 
     def _trigger_delays(self, labelled=True, by_pulse=False,
                         pulse_dim='pulseId'):

src/extra/components/timepix.py

@@ -10,6 +10,7 @@
 
 from extra_data import by_id
 from extra_data.read_machinery import split_trains
+from .utils import _select_subcomponent_trains
 
 
 # Pulse indices used internally for virtual leading and trailing pulses.
@@ -117,22 +118,18 @@ def select_trains(self, trains):
         This method accepts the same type of arguments as
         [DataCollection.select_trains][extra_data.DataCollection.select_trains].
         """
-        res = copy(self)
-
-        if self._raw_control_src is not None:
-            res._raw_control_src = self._raw_control_src.select_trains(trains)
-        if self._raw_instrument_src is not None:
-            res._raw_instrument_src = self._raw_instrument_src.select_trains(trains)
+        res = _select_subcomponent_trains(self, [
+            '_raw_control_src',
+            '_raw_instrument_src',
+            '_centroids_control_src',
+            '_centroids_instrument_src',
+            '_pulses',
+        ], trains)
         if self._raw_instrument_dc is not None:
             raw_train_ids = res.raw_size_key.drop_empty_trains().train_ids
             res._raw_instrument_dc = self._raw_instrument_dc.select_trains(
                 by_id[raw_train_ids]
             )
-        if self._centroids_control_src is not None:
-            res._centroids_control_src = self._centroids_control_src.select_trains(trains)
-        if self._centroids_instrument_src is not None:
-            res._centroids_instrument_src = self._centroids_instrument_src.select_trains(trains)
-        res._pulses = self._pulses.select_trains(trains)
 
         return res
 

src/extra/components/utils.py

@@ -63,13 +63,13 @@ def _instrument_to_sase(instrument):
         return 3
 
 
-def _select_subcomponent_trains(src, keys, dst=None):
+def _select_subcomponent_trains(src, keys, trains, dst=None):
     if dst is None:
         from copy import copy
         dst = copy(src)
 
     for key in keys:
-        prop = getattr(self, key)
+        prop = getattr(src, key)
 
         if prop is not None:
             setattr(dst, key, prop.select_trains(trains))