From 333db5c8f935aa8a92d3cc6f033c436fd5a01fba Mon Sep 17 00:00:00 2001
From: Jan Grewe <jan.grewe@g-node.org>
Date: Mon, 3 Aug 2020 18:21:36 +0200
Subject: [PATCH] [relacs] adaptations of find methods

---
 fishbook/frontend/relacs_classes.py | 20 +++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/fishbook/frontend/relacs_classes.py b/fishbook/frontend/relacs_classes.py
index 559942c..6a441c2 100644
--- a/fishbook/frontend/relacs_classes.py
+++ b/fishbook/frontend/relacs_classes.py
@@ -27,7 +27,7 @@ class BaselineData:
     def _get_data(self):
         if not self.__dataset:
             return
-        self.__repros = RePro.find("BaselineActivity", cell_id=self.__cell.id)
+        self.__repros, _ = RePro.find("BaselineActivity", cell_id=self.__cell.id)
         for i in tqdm(range(len(self.__repros)), desc="loading data"):
             r = self.__repros[i]
             sd = self.__read_spike_data(r)
@@ -311,7 +311,7 @@ class FIData:
     def _get_data(self):
         if not self.__dataset:
             return
-        self.__repros = RePro.find("FICurve", cell_id=self.__cell.id)
+        self.__repros,_  = RePro.find("FICurve", cell_id=self.__cell.id)
         for r in self.__repros:
             sd, c, eods, time = self.__read_spike_data(r)
             if sd is not None and len(sd) > 1:
@@ -365,8 +365,8 @@ class FIData:
         eods = []
         time = []
         contrasts = []
-        stimuli = Stimulus.find(cell_id=repro.cell_id, repro_id=repro.id)
-        if len(stimuli) == 0:
+        stimuli, count = Stimulus.find(cell_id=repro.cell_id, repro_id=repro.id)
+        if count == 0:
             return spikes, contrasts, eods, time
         data_source = os.path.join(self.__dataset.data_source, self.__dataset.id + ".nix")
         if not os.path.exists(data_source):
@@ -376,7 +376,7 @@ class FIData:
         b = f.blocks[0]
         self.__all_spikes = b.data_arrays["Spikes-1"][:]
         mt = None
-        for i in tqdm(range(len(stimuli)), desc="Loading data"):
+        for i in tqdm(range(count), desc="Loading data"):
             s = stimuli[i]
             if not mt or mt.id != s.multi_tag_id:
                 mt = b.multi_tags[s.multi_tag_id]
@@ -539,12 +539,18 @@ class FIData:
 
 
 class FileStimulusData:
+    """The FileStimulus class provides access to the data recorded and the stimulus presented (if accessible) 
+    during runs of the FileStimulus repro. Since the FileStimulus repro can put out any stimulus this class does not 
+    provide any further analyses. 
 
+    As any other relacs class it is instantiated with a Dataset entity.
+    """
     def __init__(self, dataset: Dataset):
         """
         Constructor.
 
-        :param dataset: The dataset entity for which the filestimulus repro data should be loaded.
+        Args
+            fishbook.Dataset: The dataset entity for which the filestimulus repro data should be loaded.
         """
         self.__spike_data = []
         self.__contrasts = []
@@ -557,7 +563,7 @@ class FileStimulusData:
     def _get_data(self):
         if not self.__dataset:
             return
-        self.__repros = RePro.find("FileStimulus", cell_id=self.__cell.id)
+        self.__repros, _ = RePro.find("FileStimulus", cell_id=self.__cell.id)
         for r in self.__repros:
             sd, c, stims = self.__read_spike_data_from_nix(r) if self.__dataset.has_nix else self.__read_spike_data_from_directory(r)
             if sd is not None and len(sd) > 1: