From 348d3a8493318b2b61d979a7f2661239b40ffbae Mon Sep 17 00:00:00 2001
From: Darshan Mandge <8694327+darshanmandge@users.noreply.github.com>
Date: Fri, 14 Jun 2024 12:33:14 +0200
Subject: [PATCH] allow stimulus_description in ScalaNWBReader (#183)

* allow stimulus_description in ScalaNWBReader

* lint fix

* edge case with no stimulus_description attribute[F

* adding AttributeError

* KeyError in ScalaNWBReader
---
 bluepyefe/nwbreader.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/bluepyefe/nwbreader.py b/bluepyefe/nwbreader.py
index 0f2bd8b..f6cb932 100644
--- a/bluepyefe/nwbreader.py
+++ b/bluepyefe/nwbreader.py
@@ -114,7 +114,14 @@ def read(self):
 
         for sweep in list(self.content['acquisition'].keys()):
             key_current = sweep.replace('Series', 'StimulusSeries')
-            protocol_name = "Step"
+            try:
+                protocol_name = self.content["acquisition"][sweep].attrs["stimulus_description"]
+            except KeyError:
+                logger.warning(f'Could not find "stimulus_description" attribute for {sweep}, Setting it as "Step"')
+                protocol_name = "Step"
+
+            if ("na" == protocol_name.lower()) or ("step" in protocol_name.lower()):
+                protocol_name = "Step"
 
             if (
                 self.target_protocols and