Add BLOB support for dict of simple types

DevDockerfile

@@ -0,0 +1,7 @@
+FROM eywalker/jupyter
+
+MAINTAINER Edgar Y. Walker <edgar.walker@gmail.com>
+
+ADD . /src
+RUN pip install -e /src &&\
+    pip install nose

datajoint/blob.py

@@ -3,7 +3,7 @@
 """
 
 import zlib
-from collections import OrderedDict
+from collections import OrderedDict, Mapping, Iterable
 import numpy as np
 from . import DataJointError
 
@@ -36,10 +36,11 @@
 
 
 class BlobReader:
-    def __init__(self, blob, simplify=False):
+    def __init__(self, blob, simplify=False, as_dict=False):
         self._simplify = simplify
         self._blob = blob
         self._pos = 0
+        self._as_dict = as_dict
 
     @property
     def pos(self):
@@ -92,10 +93,10 @@ def read_array(self, advance=True, n_bytes=None):
         dtype = dtype_list[dtype_id]
         is_complex = self.read_value('uint32')
 
-        if dtype_id == 4: # if dealing with character array
+        if dtype_id == 4:  # if dealing with character array
             data = self.read_value(dtype, count=2 * n_elem)
-            data = data[::2].astype('<U1')
-            if n_dims == 2 and shape[0] == 1:
+            data = data[::2].astype('U1')
+            if n_dims == 2 and shape[0] == 1 or n_dims == 1:
                 compact = data.squeeze()
                 data = compact if compact.shape == () else np.array(''.join(data.squeeze()))
                 shape = (1,)
@@ -129,17 +130,22 @@ def read_structure(self, advance=True, n_bytes=None):
         dt = [(f, np.object) for f in field_names]
         raw_data = []
         for k in range(n_elem):
-            vals = []
+            values = []
             for i in range(n_field):
                 nb = int(self.read_value('uint64')) # dealing with a weird bug of numpy
-                vals.append(self.read_mym_data(n_bytes=nb))
-            raw_data.append(tuple(vals))
-        data = np.rec.array(raw_data, dtype=dt)
+                values.append(self.read_mym_data(n_bytes=nb))
+            raw_data.append(tuple(values))
         if n_bytes is not None:
             assert self.pos - start == n_bytes
         if not advance:
             self.pos = start
-        return self.simplify(data.reshape(shape, order='F'))
+
+        if self._as_dict and n_elem == 1:
+            data = dict(zip(field_names, values))
+            return data
+        else:
+            data = np.rec.array(raw_data, dtype=dt)
+            return self.simplify(data.reshape(shape, order='F'))
 
     def simplify(self, array):
         """
@@ -201,41 +207,95 @@ def __str__(self):
         return str(self._blob[self.pos:])
 
 
-def pack(obj):
-    """
-    Packs an object into a blob to be compatible with mym.mex
+def pack(obj, compress=True):
+    blob = b"mYm\0"
+    blob += pack_obj(obj)
 
-    :param obj: object to be packed
-    :type obj: numpy.ndarray
-    """
-    if not isinstance(obj, np.ndarray):
-        raise DataJointError("Only numpy arrays can be saved in blobs")
+    if compress:
+        compressed = b'ZL123\0' + np.uint64(len(blob)).tostring() + zlib.compress(blob)
+        if len(compressed) < len(blob):
+            blob = compressed
+
+    return blob
+
+
+def pack_obj(obj):
+    blob = b''
+    if isinstance(obj, np.ndarray):
+        blob += pack_array(obj)
+    elif isinstance(obj, Mapping):  # TODO: check if this is a good inheritance check for dict etc.
+        blob += pack_dict(obj)
+    elif isinstance(obj, str):
+        blob += pack_array(np.array(obj, dtype=np.dtype('c')))
+    elif isinstance(obj, Iterable):
+        blob += pack_array(np.array(list(obj)))
+    elif isinstance(obj, int) or isinstance(obj, float):
+        blob += pack_array(np.array(obj))
+    else:
+        raise DataJointError("Packing object of type %s currently not supported!" % type(obj))
+
+    return blob
+
+
+def pack_array(array):
+    if not isinstance(array, np.ndarray):
+        raise ValueError("argument must be a numpy array!")
 
-    blob = b"mYm\0A"  # TODO: extend to process other data types besides arrays
-    blob += np.asarray((len(obj.shape),) + obj.shape, dtype=np.uint64).tostring()
+    blob = b"A"
+    blob += np.array((len(array.shape), ) + array.shape, dtype=np.uint64).tostring()
 
-    is_complex = np.iscomplexobj(obj)
+    is_complex = np.iscomplexobj(array)
     if is_complex:
-        obj, imaginary = np.real(obj), np.imag(obj)
+        array, imaginary = np.real(array), np.imag(array)
 
-    type_number = rev_class_id[obj.dtype]
-    assert dtype_list[type_number] == obj.dtype, 'ambiguous or unknown array type'
-    blob += np.asarray(type_number, dtype=np.uint32).tostring()
-    blob += np.int8(is_complex).tostring() + b'\0\0\0'
-    blob += obj.tostring(order='F')
+    type_number = rev_class_id[array.dtype]
+
+    if dtype_list[type_number] is None:
+        raise DataJointError("Type %s is ambiguous or unknown" % array.dtype)
+
+    blob += np.array(type_number, dtype=np.uint32).tostring()
+
+    blob += np.int32(is_complex).tostring()
+    if type_number == 4:  # if dealing with character array
+        blob += ('\x00'.join(array.tostring(order='F').decode()) + '\x00').encode()
+    else:
+        blob += array.tostring(order='F')
 
     if is_complex:
         blob += imaginary.tostring(order='F')
 
-    compressed = b'ZL123\0' + np.uint64(len(blob)).tostring() + zlib.compress(blob)
-    if len(compressed) < len(blob):
-        blob = compressed
     return blob
 
 
-def unpack(blob):
+def pack_string(value):
+    return value.encode('ascii') + b'\0'
+
+
+def pack_dict(obj):
+    """
+    Write dictionary object as a singular structure array
+    :param obj: dictionary object to serialize. The fields must be simple scalar or an array.
+    """
+    obj = OrderedDict(obj)
+    blob = b'S'
+    blob += np.array((2, 1, 1), dtype=np.uint64).tostring()
+    blob += np.array(len(obj), dtype=np.uint32).tostring()
+
+    # write out field names
+    for k in obj:
+        blob += pack_string(k)
+
+    for k, v in obj.items():
+        blob_part = pack_obj(v)
+        blob += np.array(len(blob_part), dtype=np.uint64).tostring()
+        blob += blob_part
+
+    return blob
+
+
+def unpack(blob, **kwargs):
     if blob is None:
         return None
 
-    return BlobReader(blob).unpack()
+    return BlobReader(blob, **kwargs).unpack()
 

docker-compose.yml

@@ -3,11 +3,13 @@ services:
   datajoint:
     build:
       context: .
-      dockerfile: JupyterDockerfile
+      dockerfile: DevDockerfile
     environment:
       - DJ_HOST=db
       - DJ_USER=root
       - DJ_PASS=simple
+    volumes:
+      - .:/src
     links:
       - db
     ports:

tests/test_blob.py

@@ -1,8 +1,7 @@
-
-
 import numpy as np
 from datajoint.blob import pack, unpack
 from numpy.testing import assert_array_equal, raises
+from nose.tools import assert_equal, assert_true
 
 
 def test_pack():
@@ -18,6 +17,15 @@ def test_pack():
     x = np.int16(np.random.randn(1, 2, 3))
     assert_array_equal(x, unpack(pack(x)), "Arrays do not match!")
 
+    x = {'name': 'Anonymous', 'age': 15}
+    assert_true(x == unpack(pack(x), as_dict=True), "Dict do not match!")
+
+    x = [1, 2, 3, 4]
+    assert_array_equal(x, unpack(pack(x)), "List did not pack/unpack correctly")
+
+    x = [1, 2, 3, 4]
+    assert_array_equal(x, unpack(pack(x.__iter__())), "Iterator did not pack/unpack correctly")
+
 
 def test_complex():
     z = np.random.randn(8, 10) + 1j*np.random.randn(8,10)