m-labs · Apr 24, 2016 · Apr 24, 2016 · Apr 24, 2016 · Apr 24, 2016 · Apr 24, 2016
Showing with 39 additions and 31 deletions.

+16 −25 artiq/frontend/artiq_influxdb.py

+8 −4 artiq/protocols/pyon.py

+15 −2 artiq/test/test_serialization.py
diff --git a/artiq/frontend/artiq_influxdb.py b/artiq/frontend/artiq_influxdb.py
@@ -6,6 +6,7 @@
 import atexit
 import fnmatch
 from functools import partial
+import time
 
 import numpy as np
 import aiohttp
@@ -63,24 +64,16 @@ def get_argparser():
     return parser
 
 
-def influxdb_str(s):
-    return '"' + s.replace('"', '\\"') + '"'
-
-
 def format_influxdb(v):
-    if isinstance(v, bool):
-        if v:
-            return "bool", "t"
-        else:
-            return "bool", "f"
-    elif np.issubdtype(type(v), int):
-        return "int", "{}i".format(v)
-    elif np.issubdtype(type(v), float):
-        return "float", "{}".format(v)
-    elif isinstance(v, str):
-        return "str", influxdb_str(v)
-    else:
-        return "pyon", influxdb_str(pyon.encode(v))
+    if np.issubdtype(type(v), np.bool_):
+        return "bool={}".format(v)
+    if np.issubdtype(type(v), np.integer):
+        return "int={}i".format(v)
+    if np.issubdtype(type(v), np.floating):
+        return "float={}".format(v)
+    if np.issubdtype(type(v), np.str_):
+        return "str=\"{}\"".format(v.replace('"', '\\"'))
+    return "pyon=\"{}\"".format(pyon.encode(v).replace('"', '\\"'))
 
 
 class DBWriter(TaskObject):
@@ -95,19 +88,19 @@ def __init__(self, base_url, user, password, database, table):
 
     def update(self, k, v):
         try:
-            self._queue.put_nowait((k, v))
+            self._queue.put_nowait((k, v, time.time()))
         except asyncio.QueueFull:
             logger.warning("failed to update dataset '%s': "
                            "too many pending updates", k)
 
     async def _do(self):
         while True:
-            k, v = await self._queue.get()
+            k, v, t = await self._queue.get()
             url = self.base_url + "/write"
             params = {"u": self.user, "p": self.password, "db": self.database,
-                      "consistency": "any", "precision": "n"}
-            fmt_ty, fmt_v = format_influxdb(v)
-            data = "{},dataset={} {}={}".format(self.table, k, fmt_ty, fmt_v)
+                      "precision": "ms"}
+            data = "{},dataset={} {} {}".format(
+                self.table, k, format_influxdb(v), round(t*1e3))
             try:
                 response = await aiohttp.request(
                     "POST", url, params=params, data=data)
@@ -116,9 +109,7 @@ async def _do(self):
                                k, exc_info=True)
             else:
                 if response.status not in (200, 204):
-                    content = (await response.content.read()).decode()
-                    if content:
-                        content = content[:-1]  # drop \n
+                    content = (await response.content.read()).decode().strip()
                     logger.warning("got HTTP status %d "
                                    "trying to update '%s': %s",
                                    response.status, k, content)

diff --git a/artiq/protocols/pyon.py b/artiq/protocols/pyon.py
@@ -5,7 +5,7 @@
 * Human-readable format compatible with the Python syntax.
 * Each object is serialized on a single line, with only ASCII characters.
 * Supports all basic Python data structures: None, booleans, integers,
-  floats, strings, tuples, lists, dictionaries.
+  floats, complex numbers, strings, tuples, lists, dictionaries.
 * Those data types are accurately reconstructed (unlike JSON where e.g. tuples
   become lists, and dictionary keys are turned into strings).
 * Supports Numpy arrays.
@@ -33,6 +33,7 @@
     bool: "bool",
     int: "number",
     float: "number",
+    complex: "number",
     str: "str",
     bytes: "bytes",
     tuple: "tuple",
@@ -46,8 +47,10 @@
 }
 
 _numpy_scalar = {
-    "int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64",
-    "float16", "float32", "float64"
+    "int8", "int16", "int32", "int64",
+    "uint8", "uint16", "uint32", "uint64",
+    "float16", "float32", "float64", "float128",
+    "complex64", "complex128", "complex256",
 }
 
 
@@ -132,7 +135,7 @@ def encode_fraction(self, x):
                                          self.encode(x.denominator))
 
     def encode_ordereddict(self, x):
-        return "OrderedDict("+ self.encode(list(x.items())) +")"
+        return "OrderedDict(" + self.encode(list(x.items())) + ")"
 
     def encode_nparray(self, x):
         r = "nparray("
@@ -186,6 +189,7 @@ def _npscalar(ty, data):
     "npscalar": _npscalar
 }
 
+
 def decode(s):
     """Parses a string in the Python syntax, reconstructs the corresponding
     object, and returns it."""

diff --git a/artiq/test/test_serialization.py b/artiq/test/test_serialization.py
@@ -14,14 +14,27 @@
     "a": np.int8(9), "b": np.int16(-98), "c": np.int32(42), "d": np.int64(-5),
     "e": np.uint8(8), "f": np.uint16(5), "g": np.uint32(4), "h": np.uint64(9),
     "x": np.float16(9.0), "y": np.float32(9.0), "z": np.float64(9.0),
+    1j: 1-9j,
+    "q": np.complex128(1j),
 }
 
 
 class PYON(unittest.TestCase):
     def test_encdec(self):
         for enc in pyon.encode, lambda x: pyon.encode(x, True):
-            self.assertEqual(pyon.decode(enc(_pyon_test_object)),
-                             _pyon_test_object)
+            with self.subTest(enc=enc):
+                self.assertEqual(pyon.decode(enc(_pyon_test_object)),
+                                 _pyon_test_object)
+
+    def test_encdec_array(self):
+        orig = {k: (np.array(v), np.array([v]))
+                for k, v in _pyon_test_object.items()
+                if np.isscalar(v)}
+        for enc in pyon.encode, lambda x: pyon.encode(x, True):
+            result = pyon.decode(enc(orig))
+            for k in orig:
+                with self.subTest(enc=enc, k=k, v=orig[k]):
+                    np.testing.assert_equal(result[k], orig[k])
 
 
 _json_test_object = {