From 12c03859582e1f64a9b0751528eb987c0e30a553 Mon Sep 17 00:00:00 2001
From: yuk <yuk@m-labs.hk>
Date: Mon, 17 Oct 2022 11:31:13 +0800
Subject: [PATCH] hdf5: Allow saving custom attributes

Signed-off-by: yuk <yuk@m-labs.hk>
---
 artiq/language/environment.py | 15 +++++++++++++++
 artiq/master/worker_db.py     | 23 +++++++++++++++++++++++
 artiq/master/worker_impl.py   |  1 +
 3 files changed, 39 insertions(+)

diff --git a/artiq/language/environment.py b/artiq/language/environment.py
index 9ce7034cce..ec44022391 100644
--- a/artiq/language/environment.py
+++ b/artiq/language/environment.py
@@ -383,6 +383,21 @@ def append_to_dataset(self, key, value):
         efficiently as incremental modifications in broadcast mode."""
         self.__dataset_mgr.append_to(key, value)
 
+    @rpc(flags={"async"})
+    def set_dataset_metadata(self, key, metadata):
+        """Attach metadata to the dataset itself, or to a key in the dataset.
+
+        The metadata is saved as HDF5 attributes if there was a call to
+        ``set_dataset(..., archive=True)`` with the same key.
+
+        :param key: If ``None``, attach the metadata to the dataset itself
+            (to the ``datasets`` group when saving to HDF5).
+            Otherwise, attach the metadata to the specified key.
+        :param metadata: Dict of key-value pair.
+            Must be compatible with HDF5 attributes.
+        """
+        self.__dataset_mgr.set_metadata(key, metadata)
+
     def get_dataset(self, key, default=NoDefault, archive=True):
         """Returns the contents of a dataset.
 
diff --git a/artiq/master/worker_db.py b/artiq/master/worker_db.py
index fccdc1c118..f52162af22 100644
--- a/artiq/master/worker_db.py
+++ b/artiq/master/worker_db.py
@@ -110,6 +110,7 @@ class DatasetManager:
     def __init__(self, ddb):
         self._broadcaster = Notifier(dict())
         self.local = dict()
+        self.hdf5_attributes = dict()
         self.archive = dict()
 
         self.ddb = ddb
@@ -163,6 +164,12 @@ def get(self, key, archive=False):
             self.archive[key] = data
         return data
 
+    def set_metadata(self, key, metadata):
+        if key:
+            self.hdf5_attributes["datasets/" + key] = metadata
+        else:
+            self.hdf5_attributes["datasets"] = metadata
+
     def write_hdf5(self, f):
         datasets_group = f.create_group("datasets")
         for k, v in self.local.items():
@@ -172,6 +179,11 @@ def write_hdf5(self, f):
         for k, v in self.archive.items():
             _write(archive_group, k, v)
 
+    def write_hdf5_attributes(self, f):
+        for k, attrs in self.hdf5_attributes.items():
+            if k in f:
+                _write_attributes(f, k, attrs)
+
 
 def _write(group, k, v):
     # Add context to exception message when the user writes a dataset that is
@@ -181,3 +193,14 @@ def _write(group, k, v):
     except TypeError as e:
         raise TypeError("Error writing dataset '{}' of type '{}': {}".format(
             k, type(v), e))
+
+
+def _write_attributes(f, k, attrs):
+    # Add context to exception message when the user writes a attribute that is
+    # not representable in HDF5.
+    try:
+        for attr_k, attr_v in attrs.items():
+            f[k].attrs[attr_k] = attr_v
+    except TypeError as e:
+        raise TypeError("Error writing attribute '{}' of type '{}': {}".format(
+            attr_k, type(attr_v), e))
diff --git a/artiq/master/worker_impl.py b/artiq/master/worker_impl.py
index 33d34ddf8b..24ff192404 100644
--- a/artiq/master/worker_impl.py
+++ b/artiq/master/worker_impl.py
@@ -289,6 +289,7 @@ def write_results():
         filename = "{:09}-{}.h5".format(rid, exp.__name__)
         with h5py.File(filename, "w") as f:
             dataset_mgr.write_hdf5(f)
+            dataset_mgr.write_hdf5_attributes(f)
             f["artiq_version"] = artiq_version
             f["rid"] = rid
             f["start_time"] = start_time