FIX/CLN: streamdevice annotation serialization

pcdshub · Sep 28, 2023 · 629e80d · 629e80d
1 parent bf316ec
commit 629e80d
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 29 deletions.
diff --git a/whatrecord/shell.py b/whatrecord/shell.py
@@ -606,7 +606,7 @@ def handle_dbLoadRecords(self, filename: str, macros: str = ""):
             if annotation:
                 key = StringWithContext(
                     self.streamdevice.metadata_key,
-                    context=record.context,
+                    context=(),
                 )
                 record.metadata[key] = annotation
 

diff --git a/whatrecord/streamdevice.py b/whatrecord/streamdevice.py
@@ -9,7 +9,7 @@
 from epicsmacrolib import split_iocsh_line
 
 from . import transformer
-from .common import (AnyPath, FullLoadContext, ShellStateHandler,
+from .common import (AnyPath, FullLoadContext, LoadContext, ShellStateHandler,
                      StringWithContext)
 from .db import PVAFieldReference, RecordInstance
 from .transformer import context_from_token
@@ -165,10 +165,10 @@ def standalone_variable_ref(self, ref, *_):
             arguments=[],
         )
 
-    def config_set(self, command):
+    def config_set(self, command: lark.Tree):
         return ConfigurationSetting(
-            name=command.data,
-            value=" ".join(command.children),
+            name=str(command.data),
+            value=" ".join(str(child) for child in command.children),
         )
 
     value_part = transformer.pass_through
@@ -337,34 +337,35 @@ def annotate_record(self, record: RecordInstance) -> Optional[Dict[str, Any]]:
             }
 
         info_field = info_field.value.strip()
-        results = {}
+        annotation: Dict[str, Any] = {
+            "context": [],
+        }
         try:
             proto_file, proto_name, *proto_args = split_iocsh_line(info_field).argv
             proto_file = proto_file.lstrip("@ ")
         except Exception:
-            results["error"] = (
+            annotation["error"] = (
                 f"Invalid StreamDevice input/output field: {info_field!r}"
             )
-            proto_file = None
-            proto_name = None
-            proto_args = []
+            return annotation
+
+        try:
+            protocol = self.load_streamdevice_protocol(proto_file)
+        except Exception as ex:
+            annotation["error"] = f"{ex.__class__.__name__}: {ex}"
+            protocol = None
+            return annotation
+
+        if proto_name in protocol.protocols:
+            specific_protocol = protocol.protocols[proto_name]
+            annotation["context"] = specific_protocol.context
+            annotation["protocol"] = specific_protocol
         else:
-            try:
-                protocol = self.load_streamdevice_protocol(proto_file)
-            except Exception as ex:
-                results["error"] = f"{ex.__class__.__name__}: {ex}"
-            else:
-                if proto_name in protocol.protocols:
-                    results["protocol"] = protocol.protocols[proto_name]
-                else:
-                    results["error"] = (
-                        f"Unknown protocol {proto_name!r} in {proto_file}; "
-                        f"options are: {list(protocol.protocols)}"
-                    )
-
-        return {
-            "protocol_file": proto_file,
-            "protocol_name": proto_name,
-            "protocol_args": proto_args,
-            **results,
-        }
+            annotation["context"] = [LoadContext(proto_file, 0)]
+            annotation["error"] = (
+                f"Unknown protocol {proto_name!r} in {proto_file}; "
+                f"options are: {list(protocol.protocols)}"
+            )
+        annotation["protocol_name"] = proto_name
+        annotation["protocol_args"] = proto_args
+        return annotation