Merge pull request #1153 from INM-6/enh/phyrawio_read_multicolumn_csvs

Read multi-column csvs in PhyIO

Author: JuliaSprenger

neo/rawio/phyrawio.py

@@ -149,17 +149,17 @@ def _parse_header(self):
 
             # Loop over list of list of dict and annotate each st
             for annotation_list in annotation_lists:
-                clust_key, property_name = tuple(annotation_list[0].
-                                                 keys())
-                if property_name == 'KSLabel':
-                    annotation_name = 'quality'
-                else:
-                    annotation_name = property_name.lower()
-                for annotation_dict in annotation_list:
-                    if int(annotation_dict[clust_key]) == clust_id:
-                        spiketrain_an[annotation_name] = \
-                            annotation_dict[property_name]
-                        break
+                clust_key, *property_names = tuple(annotation_list[0].keys())
+                for property_name in property_names:
+                    if property_name == 'KSLabel':
+                        annotation_name = 'quality'
+                    else:
+                        annotation_name = property_name.lower()
+                    for annotation_dict in annotation_list:
+                        if int(annotation_dict[clust_key]) == clust_id:
+                            spiketrain_an[annotation_name] = \
+                                annotation_dict[property_name]
+                            break
 
             cluster_mask = (self._spike_clusters == clust_id).flatten()
 
@@ -256,7 +256,7 @@ def _rescale_epoch_duration(self, raw_duration, dtype):
     def _parse_tsv_or_csv_to_list_of_dict(filename):
         list_of_dict = list()
         letter_pattern = re.compile('[a-zA-Z]')
-        float_pattern = re.compile(r'\d*\.')
+        float_pattern = re.compile(r'-?\d*\.')
         with open(filename) as csvfile:
             if filename.suffix == '.csv':
                 reader = csv.DictReader(csvfile, delimiter=',')
@@ -268,15 +268,19 @@ def _parse_tsv_or_csv_to_list_of_dict(filename):
 
             for row in reader:
                 if line == 0:
-                    key1, key2 = tuple(row.keys())
+                    cluster_id_key, *annotation_keys = tuple(row.keys())
                 # Convert cluster ID to int
-                row[key1] = int(row[key1])
+                row[cluster_id_key] = int(row[cluster_id_key])
                 # Convert strings without letters
-                if letter_pattern.match(row[key2]) is None:
-                    if float_pattern.match(row[key2]) is None:
-                        row[key2] = int(row[key2])
-                    else:
-                        row[key2] = float(row[key2])
+                for key in annotation_keys:
+                    value = row[key]
+                    if not len(value):
+                        row[key] = None
+                    elif letter_pattern.match(value) is None:
+                        if float_pattern.match(value) is None:
+                            row[key] = int(value)
+                        else:
+                            row[key] = float(value)
 
                 list_of_dict.append(row)
                 line += 1

neo/test/rawiotest/test_phyrawio.py

@@ -31,26 +31,35 @@ def test_csv_tsv_parser_with_csv(self):
         csv_tempfile = Path(tempfile.gettempdir()).joinpath('test.csv')
         with open(csv_tempfile, 'w') as csv_file:
             csv_writer = csv.writer(csv_file, delimiter=',')
-            csv_writer.writerow(['cluster_id', 'some_annotation'])
-            csv_writer.writerow([1, 'Good'])
-            csv_writer.writerow([2, 10])
-            csv_writer.writerow([3, 1.23])
+            csv_writer.writerow(['cluster_id', 'some_annotation', 'some_other_annotation'])
+            csv_writer.writerow([1, 'Good', 'Bad'])
+            csv_writer.writerow([2, 10, -2])
+            csv_writer.writerow([3, 1.23, -0.38])
 
         # the parser in PhyRawIO runs csv.DictReader to parse the file
         # csv.DictReader for python version 3.6+ returns list of OrderedDict
         if (3, 6) <= sys.version_info < (3, 8):
             target = [OrderedDict({'cluster_id': 1,
-                                   'some_annotation': 'Good'}),
+                                   'some_annotation': 'Good',
+                                   'some_other_annotation': 'Bad'}),
                       OrderedDict({'cluster_id': 2,
-                                   'some_annotation': 10}),
+                                   'some_annotation': 10,
+                                   'some_other_annotation': -2}),
                       OrderedDict({'cluster_id': 3,
-                                   'some_annotation': 1.23})]
+                                   'some_annotation': 1.23,
+                                   'some_other_annotation': -0.38})]
 
         # csv.DictReader for python version 3.8+ returns list of dict
         elif sys.version_info >= (3, 8):
-            target = [{'cluster_id': 1, 'some_annotation': 'Good'},
-                      {'cluster_id': 2, 'some_annotation': 10},
-                      {'cluster_id': 3, 'some_annotation': 1.23}]
+            target = [{'cluster_id': 1,
+                       'some_annotation': 'Good',
+                       'some_other_annotation': 'Bad'},
+                      {'cluster_id': 2,
+                       'some_annotation': 10,
+                       'some_other_annotation': -2},
+                      {'cluster_id': 3,
+                       'some_annotation': 1.23,
+                       'some_other_annotation': -0.38}]
 
         list_of_dict = PhyRawIO._parse_tsv_or_csv_to_list_of_dict(csv_tempfile)