Merge pull request #1440 from NxNiki/multiple_filenames

apdavison · web-flow · commit ad38785c191d · 2024-07-24T09:59:21.000+02:00
NeuralynxIO: support List[str] for filename parameter
diff --git a/neo/io/neuralynxio.py b/neo/io/neuralynxio.py
@@ -9,6 +9,7 @@
 Author: Julia Sprenger, Carlos Canova
 """
 
+import warnings
 from neo.io.basefromrawio import BaseFromRaw
 from neo.rawio.neuralynxrawio.neuralynxrawio import NeuralynxRawIO
 
@@ -29,12 +30,14 @@ class NeuralynxIO(NeuralynxRawIO, BaseFromRaw):
 
     def __init__(
         self,
-        dirname="",
-        filename="",
+        dirname,
         use_cache=False,
         cache_path="same_as_resource",
-        exclude_filename=None,
+        include_filenames=None,
+        exclude_filenames=None,
         keep_original_times=False,
+        filename=None,
+        exclude_filename=None,
     ):
         """
         Initialise IO instance
@@ -44,6 +47,7 @@ def __init__(
         dirname : str
             Directory containing data files
         filename : str
+            Deprecated and will be removed. Please use `include_filenames` instead
             Name of a single ncs, nse, nev, or ntt file to include in dataset. Will be ignored,
             if dirname is provided.
         use_cache : bool, optional
@@ -52,24 +56,39 @@ def __init__(
         cache_path : str, optional
             Folder path to use for cache files.
             Default: 'same_as_resource'
-        exclude_filename: str or list
+        exclude_filename: None,
+            Deprecated and will be removed. Please use `exclude_filenames` instead
+        include_filenames: str or list
+            Filename or list of filenames to be included. This can be absolute path or path relative to dirname.
+        exclude_filenames: str or list
             Filename or list of filenames to be excluded. Expects base filenames without
             directory path.
         keep_original_times : bool
             Preserve original time stamps as in data files. By default datasets are
             shifted to begin at t_start = 0*pq.second.
             Default: False
         """
+
+        if filename is not None:
+            warnings.warn('Deprecated and will be removed. Please use `include_filenames` instead')
+            include_filenames = [filename]
+
+        if exclude_filename is not None:
+            warnings.warn('Deprecated and will be removed. Please use `exclude_filenames` instead')
+            exclude_filenames = exclude_filename
+
         NeuralynxRawIO.__init__(
             self,
             dirname=dirname,
-            filename=filename,
-            use_cache=use_cache,
-            cache_path=cache_path,
-            exclude_filename=exclude_filename,
+            include_filenames=include_filenames,
+            exclude_filenames=exclude_filenames,
             keep_original_times=keep_original_times,
+            use_cache=use_cache,
+            cache_path=cache_path
         )
-        if self.rawmode == "one-file":
-            BaseFromRaw.__init__(self, filename)
-        elif self.rawmode == "one-dir":
+
+        if self.rawmode == "one-dir":
             BaseFromRaw.__init__(self, dirname)
+        elif self.rawmode == "multiple-files":
+            BaseFromRaw.__init__(self, include_filenames=include_filenames)
+
diff --git a/neo/rawio/neuralynxrawio/neuralynxrawio.py b/neo/rawio/neuralynxrawio/neuralynxrawio.py
@@ -57,6 +57,7 @@
 import os
 import pathlib
 import copy
+import warnings
 from collections import namedtuple, OrderedDict
 
 from neo.rawio.neuralynxrawio.ncssections import NcsSection, NcsSectionsFactory
@@ -75,10 +76,12 @@ class NeuralynxRawIO(BaseRawIO):
     dirname: str, default: ''
         Name of directory containing all files for a dataset. If provided, filename is
         ignored. But one of either dirname or filename is required.
-    filename: str, default: ''
-        Name of a single ncs, nse, nev, or ntt file to include in dataset. Will be ignored,
-        if dirname is provided. But one of either dirname or filename is required.
-    exclude_filename: str | list | None, default: None
+    include_filenames: str | list | None, default: None
+        Name of a single ncs, nse, nev or ntt file or list of such files. Will only include
+        file names in the list. This can be plain filenames or fullpath or path relative to
+        dirname.
+        All files should be in a single path.
+    exclude_filenames: str | list | None, default: None
         Name of a single ncs, nse, nev or ntt file or list of such files. Expects plain
         filenames (without directory path).
         None will search for all file types
@@ -128,32 +131,59 @@ class NeuralynxRawIO(BaseRawIO):
         ("channel_id", "uint32"),
         ("sample_rate", "uint32"),
         ("nb_valid", "uint32"),
-        ("samples", "int16", (NcsSection._RECORD_SIZE)),
+        ("samples", "int16", NcsSection._RECORD_SIZE),
     ]
 
     def __init__(
-        self, dirname="", filename="", exclude_filename=None, keep_original_times=False, strict_gap_mode=True, **kargs
+            self,
+            dirname="",
+            include_filenames=None,
+            exclude_filenames=None,
+            keep_original_times=False,
+            strict_gap_mode=True,
+            filename=None,
+            exclude_filename=None,
+            **kargs
     ):
 
-        if dirname != "":
-            self.dirname = dirname
-            self.rawmode = "one-dir"
-        elif filename != "":
-            self.filename = filename
-            self.rawmode = "one-file"
+        if not dirname:
+            raise ValueError("`dirname` cannot be empty.")
+
+        if filename is not None:
+            include_filenames = [filename]
+            warnings.warn("`filename` is deprecated and will be removed. Please use `include_filenames` instead")
+
+        if exclude_filename is not None:
+            if isinstance(exclude_filename, str):
+                exclude_filenames = [exclude_filename]
+            else:
+                exclude_filenames = exclude_filename
+            warnings.warn("`exclude_filename` is deprecated and will be removed. Please use `exclude_filenames` instead")
+
+        if include_filenames is None:
+            include_filenames = []
+        elif isinstance(include_filenames, str):
+            include_filenames = [include_filenames]
+
+        if exclude_filenames is None:
+            exclude_filenames = []
+        elif isinstance(exclude_filenames, str):
+            exclude_filenames = [exclude_filenames]
+
+        if include_filenames:
+            self.rawmode = 'multiple-files'
         else:
-            raise ValueError("One of dirname or filename must be provided.")
+            self.rawmode = "one-dir"
 
+        self.dirname = dirname
+        self.include_filenames = include_filenames
+        self.exclude_filenames = exclude_filenames
         self.keep_original_times = keep_original_times
         self.strict_gap_mode = strict_gap_mode
-        self.exclude_filename = exclude_filename
         BaseRawIO.__init__(self, **kargs)
 
     def _source_name(self):
-        if self.rawmode == "one-file":
-            return self.filename
-        else:
-            return self.dirname
+        return self.dirname
 
     def _parse_header(self):
 
@@ -183,33 +213,24 @@ def _parse_header(self):
         event_annotations = []
 
         if self.rawmode == "one-dir":
-            filenames = sorted(os.listdir(self.dirname))
-            dirname = self.dirname
+            filenames = os.listdir(self.dirname)
         else:
-            if not os.path.isfile(self.filename):
+            filenames = self.include_filenames
+
+        filenames = [f for f in filenames if f not in self.exclude_filenames]
+        full_filenames = [os.path.join(self.dirname, f) for f in filenames]
+
+        for filename in full_filenames:
+            if not os.path.isfile(filename):
                 raise ValueError(
                     f"Provided Filename is not a file: "
-                    f"{self.filename}. If you want to provide a "
+                    f"{filename}. If you want to provide a "
                     f"directory use the `dirname` keyword"
                 )
 
-            dirname, fname = os.path.split(self.filename)
-            filenames = [fname]
-
-        if not isinstance(self.exclude_filename, (list, set, np.ndarray)):
-            self.exclude_filename = [self.exclude_filename]
-
-        # remove files that were explicitly excluded
-        if self.exclude_filename is not None:
-            for excl_file in self.exclude_filename:
-                if excl_file in filenames:
-                    filenames.remove(excl_file)
-
         stream_props = {}  # {(sampling_rate, n_samples, t_start): {stream_id: [filenames]}
 
-        for filename in filenames:
-            filename = os.path.join(dirname, filename)
-
+        for filename in full_filenames:
             _, ext = os.path.splitext(filename)
             ext = ext[1:]  # remove dot
             ext = ext.lower()  # make lower case for comparisons
diff --git a/neo/test/iotest/test_neuralynxio.py b/neo/test/iotest/test_neuralynxio.py
@@ -180,20 +180,21 @@ def test_read_block(self):
         block = nio.read_block(signal_group_mode="group-by-same-units")
         self.assertEqual(len(block.groups), 1)
 
-    def test_read_single_file(self):
+    def test_include_filenames(self):
         filename = self.get_local_path("neuralynx/Cheetah_v5.7.4/original_data/CSC1.ncs")
-        nio = NeuralynxIO(filename=filename, use_cache=False)
+        dirname, filename = os.path.split(filename)
+        nio = NeuralynxIO(dirname=dirname, include_filenames=filename, use_cache=False)
         block = nio.read_block()
         self.assertTrue(len(block.segments[0].analogsignals) > 0)
         self.assertTrue((len(block.segments[0].spiketrains)) == 0)
         self.assertTrue((len(block.segments[0].events)) == 0)
         self.assertTrue((len(block.segments[0].epochs)) == 0)
 
-    def test_exclude_filename(self):
+    def test_exclude_filenames(self):
         dname = self.get_local_path("neuralynx/Cheetah_v5.7.4/original_data/")
 
         # exclude a single file
-        nio = NeuralynxIO(dirname=dname, exclude_filename="CSC1.ncs", use_cache=False)
+        nio = NeuralynxIO(dirname=dname, exclude_filenames="CSC1.ncs", use_cache=False)
         block = nio.read_block()
         self.assertTrue(len(block.segments[0].analogsignals) > 0)
         self.assertTrue((len(block.segments[0].spiketrains)) >= 0)
@@ -202,7 +203,7 @@ def test_exclude_filename(self):
 
         # exclude all ncs files from session
         exclude_files = [f"CSC{i}.ncs" for i in range(6)]
-        nio = NeuralynxIO(dirname=dname, exclude_filename=exclude_files, use_cache=False)
+        nio = NeuralynxIO(dirname=dname, exclude_filenames=exclude_files, use_cache=False)
         block = nio.read_block()
         self.assertTrue(len(block.segments[0].analogsignals) == 0)
         self.assertTrue((len(block.segments[0].spiketrains)) >= 0)
diff --git a/neo/test/rawiotest/test_neuralynxrawio.py b/neo/test/rawiotest/test_neuralynxrawio.py
@@ -1,5 +1,6 @@
 import unittest
 
+import os
 import numpy as np
 
 from neo.rawio.neuralynxrawio.neuralynxrawio import NeuralynxRawIO
@@ -105,14 +106,15 @@ def test_scan_ncs_files(self):
         # check that there are only 3 memmaps
         self.assertEqual(len(rawio._sigs_memmaps[seg_idx]), 3)
 
-    def test_single_file_mode(self):
+    def test_include_filenames(self):
         """
-        Tests reading of single files.
+        Tests include_filenames with only one file
         """
 
         # test single analog signal channel
         fname = self.get_local_path("neuralynx/Cheetah_v5.6.3/original_data/CSC1.ncs")
-        rawio = NeuralynxRawIO(filename=fname)
+        dirname, filename = os.path.split(fname)
+        rawio = NeuralynxRawIO(dirname=dirname, include_filenames=filename)
         rawio.parse_header()
 
         self.assertEqual(rawio._nb_segment, 2)
@@ -127,7 +129,8 @@ def test_single_file_mode(self):
 
         # test one single electrode channel
         fname = self.get_local_path("neuralynx/Cheetah_v5.5.1/original_data/STet3a.nse")
-        rawio = NeuralynxRawIO(filename=fname)
+        dirname, filename = os.path.split(fname)
+        rawio = NeuralynxRawIO(dirname=dirname, include_filenames=filename)
         rawio.parse_header()
 
         self.assertEqual(rawio._nb_segment, 1)
@@ -143,7 +146,7 @@ def test_single_file_mode(self):
     def test_exclude_filenames(self):
         # exclude single ncs file from session
         dname = self.get_local_path("neuralynx/Cheetah_v5.6.3/original_data/")
-        rawio = NeuralynxRawIO(dirname=dname, exclude_filename="CSC2.ncs")
+        rawio = NeuralynxRawIO(dirname=dname, exclude_filenames="CSC2.ncs")
         rawio.parse_header()
 
         self.assertEqual(rawio._nb_segment, 2)
@@ -157,7 +160,7 @@ def test_exclude_filenames(self):
         self.assertEqual(len(rawio.header["event_channels"]), 2)
 
         # exclude multiple files from session
-        rawio = NeuralynxRawIO(dirname=dname, exclude_filename=["Events.nev", "CSC2.ncs"])
+        rawio = NeuralynxRawIO(dirname=dname, exclude_filenames=["Events.nev", "CSC2.ncs"])
         rawio.parse_header()
 
         self.assertEqual(rawio._nb_segment, 2)
@@ -378,6 +381,7 @@ def test_equality(self):
     # test = TestNeuralynxRawIO()
     # test.test_scan_ncs_files()
     # test.test_exclude_filenames()
+    # test.test_include_filenames()
 
     # test = TestNcsSectionsFactory()
     # test.test_ncsblocks_partial()
