generalize DecodingParameters fetch1

LorenFrankLab · Nov 26, 2024 · 5e1d26d · 5e1d26d
1 parent e66a8af
commit 5e1d26d
Showing 1 changed file with 14 additions and 1 deletion.
diff --git a/src/spyglass/decoding/v1/core.py b/src/spyglass/decoding/v1/core.py
@@ -97,7 +97,20 @@ def fetch(self, *args, **kwargs):
     def fetch1(self, *args, **kwargs):
         """Return one decoding paramset as a class."""
         row = super().fetch1(*args, **kwargs)
-        row["decoding_params"] = restore_classes(row["decoding_params"])
+
+        if len(args) == 0:
+            row["decoding_params"] = restore_classes(row["decoding_params"])
+            return row
+
+        if "decoding_params" in args:
+            if len(args) == 1:
+                return restore_classes(row)
+            row = list(row)
+            row[args.index("decoding_params")] = restore_classes(
+                row[args.index("decoding_params")]
+            )
+            return tuple(row)
+
         return row