From fcd29492f92405c1c184267d25a469e93582adc4 Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Sat, 24 Oct 2015 10:54:59 +0800
Subject: [PATCH] worker_db: support more types in HDF5 output. Closes #144.
 Closes #121

---
 .gitignore                |  1 +
 artiq/master/worker_db.py | 34 ++++++++++++++++++++++++++--------
 artiq/test/h5types.py     | 33 +++++++++++++++++++++++++++++++++
 3 files changed, 60 insertions(+), 8 deletions(-)
 create mode 100644 artiq/test/h5types.py

diff --git a/.gitignore b/.gitignore
index 50894b653..2cc1c622f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -12,6 +12,7 @@ doc/manual/_build
 /*.egg-info
 /.coverage
 artiq/test/results
+artiq/test/h5types.h5
 examples/master/results
 examples/master/dataset_db.pyon
 examples/sim/dataset_db.pyon
diff --git a/artiq/master/worker_db.py b/artiq/master/worker_db.py
index 950ec147f..17dee5a1b 100644
--- a/artiq/master/worker_db.py
+++ b/artiq/master/worker_db.py
@@ -5,7 +5,7 @@ import os
 import time
 import re
 
-import numpy
+import numpy as np
 import h5py
 
 from artiq.protocols.sync_struct import Notifier
@@ -119,7 +119,21 @@ def get_last_rid():
 
 _type_to_hdf5 = {
     int: h5py.h5t.STD_I64BE,
-    float: h5py.h5t.IEEE_F64BE
+    float: h5py.h5t.IEEE_F64BE,
+
+    np.int8: h5py.h5t.STD_I8BE,
+    np.int16: h5py.h5t.STD_I16BE,
+    np.int32: h5py.h5t.STD_I32BE,
+    np.int64: h5py.h5t.STD_I64BE,
+
+    np.uint8: h5py.h5t.STD_U8BE,
+    np.uint16: h5py.h5t.STD_U16BE,
+    np.uint32: h5py.h5t.STD_U32BE,
+    np.uint64: h5py.h5t.STD_U64BE,
+
+    np.float16: h5py.h5t.IEEE_F16BE,
+    np.float32: h5py.h5t.IEEE_F32BE,
+    np.float64: h5py.h5t.IEEE_F64BE
 }
 
 def result_dict_to_hdf5(f, rd):
@@ -137,15 +151,19 @@ def result_dict_to_hdf5(f, rd):
                                 " HDF5 output".format(el_ty))
             dataset = f.create_dataset(name, (len(data), ), el_ty_h5)
             dataset[:] = data
-        elif isinstance(data, numpy.ndarray):
+        elif isinstance(data, np.ndarray):
             f.create_dataset(name, data=data)
         else:
             ty = type(data)
-            try:
-                ty_h5 = _type_to_hdf5[ty]
-            except KeyError:
-                raise TypeError("Type {} is not supported for HDF5 output"
-                                .format(ty))
+            if ty is str:
+                ty_h5 = "S{}".format(len(data))
+                data = data.encode()
+            else:
+                try:
+                    ty_h5 = _type_to_hdf5[ty]
+                except KeyError:
+                    raise TypeError("Type {} is not supported for HDF5 output"
+                                    .format(ty))
             dataset = f.create_dataset(name, (), ty_h5)
             dataset[()] = data
 
diff --git a/artiq/test/h5types.py b/artiq/test/h5types.py
new file mode 100644
index 000000000..dd204aa66
--- /dev/null
+++ b/artiq/test/h5types.py
@@ -0,0 +1,33 @@
+import unittest
+
+import h5py
+import numpy as np
+
+from artiq.master.worker_db import result_dict_to_hdf5
+
+
+class TypesCase(unittest.TestCase):
+    def test_types(self):
+        d = {
+            "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)]
+        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)]
+
+        with h5py.File("h5types.h5", "w") as f:
+            result_dict_to_hdf5(f, d)