Properly copy attributes of CONTROL leafs and RUN-only keys

src/exdf/write/sd_writer.py

@@ -10,7 +10,7 @@ stricter validation when combining data from different locations.
 
 
 from functools import reduce
-from itertools import accumulate
+from itertools import accumulate, product
 from logging import getLogger
 from operator import or_
 from os.path import basename
@@ -207,22 +207,35 @@ class SourceDataWriter:
                 if run_entry:
                     run_entry = (run_entry['value'], run_entry['timestamp'])
 
-                ctrl_values = sd[f'{key}.value'].ndarray()
-                ctrl_timestamps = sd[f'{key}.timestamp'].ndarray()
+                value_key = f'{key}.value'
+                timestamp_key = f'{key}.timestamp'
+
+                ctrl_values = sd[value_key].ndarray()
+                ctrl_timestamps = sd[timestamp_key].ndarray()
 
                 h5source.create_key(
                     key, values=ctrl_values, timestamps=ctrl_timestamps,
-                    run_entry=run_entry, attrs=attrs.pop(key, None))
+                    run_entry=run_entry, attrs=attrs.pop(key, None),
+                    attrs_value=attrs.pop(value_key, None),
+                    attrs_timestamp=attrs.pop(timestamp_key, None))
 
             # Write remaining RUN-only keys.
             for key, leafs in run_data_leafs.items():
                 h5source.create_run_key(
-                    key, **leafs, attrs=attrs.pop(key, None))
+                    key, **leafs, attrs=attrs.pop(key, None),
+                    attrs_value=attrs.pop(f'{key}.value', None),
+                    attrs_timestamp=attrs.pop(f'{key}.timestamp', None))
 
             # Fill in the missing attributes for nodes.
             for path, attrs in attrs.items():
                 h5source.run_key[path].attrs.update(attrs)
-                h5source.key[path].attrs.update(attrs)
+
+                try:
+                    grp = h5source.key[path]
+                except KeyError:
+                    pass  # May be RUN-only.
+                else:
+                    grp.attrs.update(attrs)
 
     def write_instrument(self, f, sources):
         """Write INSTRUMENT data.
@@ -436,22 +449,39 @@ def build_sources_index(sources):
 
 
 def get_key_attributes(sd):
+    if sd.is_control:
+        section = 'RUN'
+        keys = {key[:key.rfind('.')] for key in sd.run_values().keys()}
+        leaf_suffixes = ['', '.value', '.timestamp']
+    else:
+        section = 'INSTRUMENT'
+        keys = sd.keys()
+        leaf_suffixes = ['']
+
     source_attrs = dict()
 
     def build_path(a, b):
         return f'{a}.{b}'
 
-    for key in sd.keys(inc_timestamps=False):
+    # Find the true files belonging to this source, assuming they're all
+    # in the same file via a random key. Importantly this will be a try
+    # CONTROL key in the case of CONTROL sources, not a RUN-only key.
+    files = [FileAccess(f) for f in sd[sd.one_key()].source_file_paths]
+
+    for key in keys:
         paths = accumulate(key.split('.'), func=build_path)
 
-        for path in paths:
-            if path in source_attrs:
+        for base_path in paths:
+            if base_path in source_attrs:
                 # Skip this path, already parent of some other key.
                 continue
 
-            for source_file in sd[key].source_file_paths:
-                fa = FileAccess(source_file)
-                hdf_path = f'{sd.section}/{sd.source}/{path.replace(".", "/")}'
+            for fa, suffix in product(
+                files, leaf_suffixes if base_path == key else ['']
+            ):
+                path = base_path + suffix
+
+                hdf_path = f'{section}/{sd.source}/{path.replace(".", "/")}'
                 path_attrs = dict(fa.file[hdf_path].attrs)
 
                 existing_attrs = source_attrs.setdefault(path, path_attrs)