add indexing of multiple metadata columns in first index, add test fo…

…r all objects

pynapple/core/interval_set.py

@@ -274,6 +274,7 @@ def __setitem__(self, key, value):
 
     def __getitem__(self, key):
         if isinstance(key, str):
+            # self[str]
             if key == "start":
                 return self.values[:, 0]
             elif key == "end":
@@ -284,11 +285,22 @@ def __getitem__(self, key):
                 raise IndexError(
                     f"Unknown string argument. Should be in {['start', 'end'] + list(self._metadata.keys())}"
                 )
+        elif isinstance(key, list) and all(isinstance(x, str) for x in key):
+            # self[[*str]]
+            # easiest to convert to dataframe and then slice
+            # in case of mixing ["start", "end"] with metadata columns
+            df = self.as_dataframe()
+            if all(x in key[1] for x in ["start", "end"]):
+                return IntervalSet(df[key])
+            else:
+                return df[key]
         elif isinstance(key, Number):
+            # self[Number]
             output = self.values.__getitem__(key)
             metadata = _MetadataBase.__getitem__(self, key)
             return IntervalSet(start=output[0], end=output[1], **metadata)
         elif isinstance(key, (slice, list, np.ndarray, pd.Series)):
+            # self[array_like]
             output = self.values.__getitem__(key)
             metadata = _MetadataBase.__getitem__(self, key).reset_index(drop=True)
             return IntervalSet(start=output[:, 0], end=output[:, 1], **metadata)

pynapple/core/time_series.py

@@ -1068,7 +1068,10 @@ def __getitem__(self, key, *args, **kwargs):
             or hasattr(key, "__iter__")
             and all([isinstance(k, str) for k in key])
         ):
-            return self.loc[key]
+            if all(k in self.metadata_columns for k in key):
+                return _MetadataBase.__getitem__(self, key)
+            else:
+                return self.loc[key]
 
         else:
             if isinstance(key, pd.Series) and key.index.equals(self.columns):

pynapple/core/ts_group.py

@@ -213,6 +213,9 @@ def __getitem__(self, key):
                 return _MetadataBase.__getitem__(self, key)
             else:
                 raise KeyError(r"Key {} not in group index.".format(key))
+        elif isinstance(key, list) and all(isinstance(k, str) for k in key):
+            # index multiple metadata columns
+            return _MetadataBase.__getitem__(self, key)
 
         # array boolean are transformed into indices
         # note that raw boolean are hashable, and won't be

tests/test_metadata.py

@@ -567,6 +567,14 @@ def test_metadata_slicing(self, obj, label, val):
         # metadata columns should be the same
         assert np.all(obj2.metadata_columns == obj.metadata_columns)
 
+    def test_metadata_index_columns(self, obj):
+        # add metadata
+        obj.set_info(one=[1, 1, 1, 1], two=[2, 2, 2, 2], three=[3, 3, 3, 3])
+
+        # test metadata columns
+        assert np.all(obj["one"] == 1)
+        assert np.all(obj[["two", "three"]] == obj._metadata[["two", "three"]])
+
     def test_save_and_load_npz(self, obj):
         obj.set_info(label1=[1, 1, 2, 2], label2=[0, 1, 2, 3])