Option to have a Tsd as an index

pynapple/core/time_series.py

@@ -104,6 +104,8 @@ def __init__(self, t, d, time_units="s", time_support=None, load_array=True):
 
     def __setitem__(self, key, value):
         """setter for time series"""
+        if isinstance(key, Tsd):
+            key = key.d
         try:
             self.values.__setitem__(key, value)
         except IndexError:
@@ -1237,14 +1239,18 @@ def __repr__(self):
                 return tabulate([], headers=headers) + "\n" + bottom
 
     def __getitem__(self, key, *args, **kwargs):
+        if isinstance(key, Tsd):
+            key = key.d
+
         output = self.values.__getitem__(key)
+
         if isinstance(key, tuple):
             index = self.index.__getitem__(key[0])
+        elif isinstance(index, Number):
+            index = np.array([index])
         else:
             index = self.index.__getitem__(key)
-
-        if isinstance(index, Number):
-            index = np.array([index])
+
 
         if all(is_array_like(a) for a in [index, output]):
             if output.shape[0] == index.shape[0]:

pynapple/io/misc.py

@@ -22,7 +22,7 @@
 from .suite2p import Suite2P
 
 
-def load_file(path):
+def load_file(path, lazy_loading=True):
     """Load file. Current format supported is (npz,nwb,)
 
     .npz -> If the file is compatible with a pynapple format, the function will return a pynapple object.
@@ -34,6 +34,8 @@ def load_file(path):
     ----------
     path : str
         Path to the file
+    lazy_loading : bool
+        Lazy loading of the data, used only for NWB files
 
     Returns
     -------
@@ -49,7 +51,7 @@ def load_file(path):
         if path.endswith(".npz"):
             return NPZFile(path).load()
         elif path.endswith(".nwb"):
-            return NWBFile(path)
+            return NWBFile(path, lazy_loading=lazy_loading)
         else:
             raise RuntimeError("File format not supported")
     else: