@@ -429,36 +429,59 @@ def compute(self, traces, peaks):
429429 return sparse_wfs
430430
431431
432- def find_parent_of_type (list_of_parents , parent_type , unique = True ):
432+ def find_parent_of_type (list_of_parents , parent_type ):
433+ """
434+ Find a single parent of a given type(s) in a list of parents.
435+ If multiple parents of the given type are found, the first parent is returned.
436+
437+ Parameters
438+ ----------
439+ list_of_parents : list of PipelineNode
440+ List of parents to search through.
441+ parent_type : type
442+ The type of parent to search for.
443+
444+ Returns
445+ -------
446+ parent : PipelineNode or None
447+ The parent of the given type. Returns None if no parent of the given type is found.
448+ """
433449 if list_of_parents is None :
434450 return None
435451
436- parents = []
437- for parent in list_of_parents :
438- if isinstance (parent , parent_type ):
439- parents .append (parent )
452+ parents = find_parents_of_type (list_of_parents , parent_type )
440453
441- if unique and len (parents ) == 1 :
442- return parents [0 ]
443- elif not unique and len (parents ) > 1 :
454+ if len (parents ) > 0 :
444455 return parents [0 ]
445456 else :
446457 return None
447458
448459
449460def find_parents_of_type (list_of_parents , parent_type ):
461+ """
462+ Find all parents of a given type(s) in a list of parents.
463+
464+ Parameters
465+ ----------
466+ list_of_parents : list of PipelineNode
467+ List of parents to search through.
468+ parent_type : type | tuple of types
469+ The type(s) of parents to search for.
470+
471+ Returns
472+ -------
473+ parents : list of PipelineNode
474+ List of parents of the given type(s). Returns an empty list if no parents of the given type(s) are found.
475+ """
450476 if list_of_parents is None :
451- return None
477+ return []
452478
453479 parents = []
454480 for parent in list_of_parents :
455481 if isinstance (parent , parent_type ):
456482 parents .append (parent )
457483
458- if len (parents ) > 0 :
459- return parents
460- else :
461- return None
484+ return parents
462485
463486
464487def check_graph (nodes ):
@@ -618,7 +641,6 @@ def _compute_peak_pipeline_chunk(segment_index, start_frame, end_frame, worker_c
618641 recording_segment = recording ._recording_segments [segment_index ]
619642 retrievers = find_parents_of_type (nodes , (SpikeRetriever , PeakRetriever ))
620643 # get peak slices once for all retrievers
621- retriever_node = None
622644 peak_slice_by_retriever = {}
623645 for retriever in retrievers :
624646 peak_slice = i0 , i1 = retriever .get_peak_slice (segment_index , start_frame , end_frame , max_margin )
0 commit comments