From 40b4129c65645c37a88cf0d6be96117150a0f870 Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Wed, 28 Oct 2015 18:33:42 +0800
Subject: [PATCH] Only support scalars and numpy arrays in HDF5 output. Update
 documentation. Closes #145

---
 artiq/language/environment.py | 10 ++++++++--
 artiq/master/worker_db.py     | 31 +++++++++++++++----------------
 artiq/test/h5types.py         | 16 ++++------------
 3 files changed, 27 insertions(+), 30 deletions(-)

diff --git a/artiq/language/environment.py b/artiq/language/environment.py
index bd7cea725..8628da8d5 100644
--- a/artiq/language/environment.py
+++ b/artiq/language/environment.py
@@ -208,9 +208,15 @@ class HasEnvironment:
                     broadcast=False, persist=False, save=True):
         """Sets the contents and handling modes of a dataset.
 
+        If the dataset is broadcasted, it must be PYON-serializable.
+        If the dataset is saved, it must be a scalar (``bool``, ``int``,
+        ``float`` or NumPy scalar) or a NumPy array.
+
         :param broadcast: the data is sent in real-time to the master, which
-            dispatches it. Returns a Notifier that can be used to mutate the dataset.
-        :param persist: the master should store the data on-disk. Implies broadcast.
+            dispatches it. Returns a Notifier that can be used to mutate the
+            dataset.
+        :param persist: the master should store the data on-disk. Implies
+            broadcast.
         :param save: the data is saved into the local storage of the current
             run (archived as a HDF5 file).
         """
diff --git a/artiq/master/worker_db.py b/artiq/master/worker_db.py
index c0cbbe73d..e490c3396 100644
--- a/artiq/master/worker_db.py
+++ b/artiq/master/worker_db.py
@@ -138,21 +138,17 @@ _type_to_hdf5 = {
 
 def result_dict_to_hdf5(f, rd):
     for name, data in rd.items():
-        if isinstance(data, list):
-            el_ty = type(data[0])
-            for d in data:
-                if type(d) != el_ty:
-                    raise TypeError("All list elements must have the same"
-                                    " type for HDF5 output")
-            try:
-                el_ty_h5 = _type_to_hdf5[el_ty]
-            except KeyError:
-                raise TypeError("List element type {} is not supported for"
-                                " HDF5 output".format(el_ty))
-            dataset = f.create_dataset(name, (len(data), ), el_ty_h5)
-            dataset[:] = data
-        elif isinstance(data, np.ndarray):
-            f.create_dataset(name, data=data)
+        flag = None
+        # beware: isinstance(True/False, int) == True
+        if isinstance(data, bool):
+            data = np.int8(data)
+            flag = "py_bool"
+        elif isinstance(data, int):
+            data = np.int64(data)
+            flag = "py_int"
+
+        if isinstance(data, np.ndarray):
+            dataset = f.create_dataset(name, data=data)
         else:
             ty = type(data)
             if ty is str:
@@ -163,10 +159,13 @@ def result_dict_to_hdf5(f, rd):
                     ty_h5 = _type_to_hdf5[ty]
                 except KeyError:
                     raise TypeError("Type {} is not supported for HDF5 output"
-                                    .format(ty))
+                                    .format(ty)) from None
             dataset = f.create_dataset(name, (), ty_h5)
             dataset[()] = data
 
+        if flag is not None:
+            dataset.attrs[flag] = np.int8(1)
+
 
 class DatasetManager:
     def __init__(self, ddb):
diff --git a/artiq/test/h5types.py b/artiq/test/h5types.py
index dd204aa66..4a3b1ca19 100644
--- a/artiq/test/h5types.py
+++ b/artiq/test/h5types.py
@@ -9,25 +9,17 @@ from artiq.master.worker_db import result_dict_to_hdf5
 class TypesCase(unittest.TestCase):
     def test_types(self):
         d = {
+            "bool": True,
             "int": 42,
             "float": 42.0,
             "string": "abcdef",
-
-            "intlist": [1, 2, 3],
-            "floatlist": [1.0, 2.0, 3.0]
         }
 
         for size in 8, 16, 32, 64:
-            signed = getattr(np, "int" + str(size))
-            unsigned = getattr(np, "uint" + str(size))
-            d["i"+str(size)] = signed(42)
-            d["u"+str(size)] = unsigned(42)
-            d["i{}list".format(size)] = [signed(x) for x in range(3)]
-            d["u{}list".format(size)] = [unsigned(x) for x in range(3)]
+            d["i"+str(size)] = getattr(np, "int" + str(size))(42)
+            d["u"+str(size)] = getattr(np, "uint" + str(size))(42)
         for size in 16, 32, 64:
-            ty = getattr(np, "float" + str(size))
-            d["f"+str(size)] = ty(42)
-            d["f{}list".format(size)] = [ty(x) for x in range(3)]
+            d["f"+str(size)] = getattr(np, "float" + str(size))(42)
 
         with h5py.File("h5types.h5", "w") as f:
             result_dict_to_hdf5(f, d)