add object dtype convenience API

Author: alimanfoo

zarr/storage.py

@@ -25,7 +25,8 @@
 
 
 from zarr.util import (normalize_shape, normalize_chunks, normalize_order,
-                       normalize_storage_path, buffer_size, normalize_fill_value, nolock)
+                       normalize_storage_path, buffer_size,
+                       normalize_fill_value, nolock, normalize_dtype)
 from zarr.meta import encode_array_metadata, encode_group_metadata
 from zarr.compat import PY2, binary_type
 from numcodecs.registry import codec_registry
@@ -308,10 +309,7 @@ def _init_array_metadata(store, shape, chunks=None, dtype=None, compressor='defa
 
     # normalize metadata
     shape = normalize_shape(shape)
-    dtype = np.dtype(dtype)
-    if dtype.kind in 'mM':
-        raise ValueError('datetime64 and timedelta64 dtypes are not currently supported; '
-                         'please store the data using int64 instead')
+    dtype, object_codec = normalize_dtype(dtype, object_codec)
     chunks = normalize_chunks(chunks, shape, dtype.itemsize)
     order = normalize_order(order)
     fill_value = normalize_fill_value(fill_value, dtype)

zarr/tests/test_core.py

@@ -21,7 +21,7 @@
                           DBMStore, LMDBStore, atexit_rmtree, atexit_rmglob)
 from zarr.core import Array
 from zarr.errors import PermissionError
-from zarr.compat import PY2
+from zarr.compat import PY2, text_type, binary_type
 from zarr.util import buffer_size
 from numcodecs import (Delta, FixedScaleOffset, Zlib, Blosc, BZ2, MsgPack, Pickle,
                        Categorize, JSON, VLenUTF8, VLenBytes, VLenArray)
@@ -941,6 +941,13 @@ def test_object_arrays_vlen_text(self):
         z[:] = data
         assert_array_equal(data, z[:])
 
+        # convenience API
+        z = self.create_array(shape=data.shape, dtype=text_type)
+        assert z.dtype == object
+        assert isinstance(z.filters[0], VLenUTF8)
+        z[:] = data
+        assert_array_equal(data, z[:])
+
         z = self.create_array(shape=data.shape, dtype=object, object_codec=MsgPack())
         z[:] = data
         assert_array_equal(data, z[:])
@@ -967,6 +974,13 @@ def test_object_arrays_vlen_bytes(self):
         z[:] = data
         assert_array_equal(data, z[:])
 
+        # convenience API
+        z = self.create_array(shape=data.shape, dtype=binary_type)
+        assert z.dtype == object
+        assert isinstance(z.filters[0], VLenBytes)
+        z[:] = data
+        assert_array_equal(data, z[:])
+
         z = self.create_array(shape=data.shape, dtype=object, object_codec=Pickle())
         z[:] = data
         assert_array_equal(data, z[:])
@@ -977,18 +991,29 @@ def test_object_arrays_vlen_array(self):
                          np.array([5]),
                          np.array([2, 8, 12])] * 1000, dtype=object)
 
+        def compare_arrays(expected, actual, item_dtype):
+            assert isinstance(actual, np.ndarray)
+            assert actual.dtype == object
+            assert actual.shape == expected.shape
+            for e, a in zip(expected.flat, actual.flat):
+                assert isinstance(a, np.ndarray)
+                assert_array_equal(e, a)
+                assert a.dtype == item_dtype
+
         codecs = VLenArray(int), VLenArray('<u4')
         for codec in codecs:
             z = self.create_array(shape=data.shape, dtype=object, object_codec=codec)
             z[:] = data
-            a = z[:]
-            assert isinstance(a, np.ndarray)
-            assert a.dtype == object
-            assert a.shape == data.shape
-            for expected, actual in zip(data.flat, a.flat):
-                assert isinstance(actual, np.ndarray)
-                assert_array_equal(expected, actual)
-                assert actual.dtype == codec.dtype
+            compare_arrays(data, z[:], codec.dtype)
+
+        # convenience API
+        for item_type in 'int', '<u4':
+            z = self.create_array(shape=data.shape, dtype='array:{}'.format(item_type))
+            assert z.dtype == object
+            assert isinstance(z.filters[0], VLenArray)
+            assert z.filters[0].dtype == np.dtype(item_type)
+            z[:] = data
+            compare_arrays(data, z[:], np.dtype(item_type))
 
     def test_object_arrays_danger(self):
 

zarr/util.py

@@ -4,16 +4,26 @@
 from textwrap import TextWrapper, dedent
 import numbers
 import uuid
+import inspect
 
 
 from asciitree import BoxStyle, LeftAligned
 from asciitree.traversal import Traversal
 import numpy as np
+from numcodecs.registry import codec_registry
 
 
 from zarr.compat import PY2, reduce, text_type, binary_type
 
 
+# codecs to use for object dtype convenience API
+object_codecs = {
+    text_type.__name__: 'vlen-utf8',
+    binary_type.__name__: 'vlen-bytes',
+    'array': 'vlen-array',
+}
+
+
 def normalize_shape(shape):
     """Convenience function to normalize the `shape` argument."""
 
@@ -116,6 +126,39 @@ def normalize_chunks(chunks, shape, typesize):
     return chunks
 
 
+def normalize_dtype(dtype, object_codec):
+
+    # convenience API for object arrays
+    if inspect.isclass(dtype):
+        dtype = dtype.__name__
+    if isinstance(dtype, str):
+        tokens = dtype.split(':')
+        key = tokens[0]
+        if key in object_codecs:
+            dtype = np.dtype(object)
+            if object_codec is None:
+                codec_id = object_codecs[key]
+                if len(tokens) > 1:
+                    args = tokens[1].split(',')
+                else:
+                    args = ()
+                try:
+                    object_codec = codec_registry[codec_id](*args)
+                except KeyError:
+                    raise ValueError('codec %r for object type %r is not '
+                                     'available; please provide an '
+                                     'object_codec manually' % (codec_id, key))
+            return dtype, object_codec
+
+    dtype = np.dtype(dtype)
+
+    if dtype.kind in 'mM':
+        raise ValueError('datetime64 and timedelta64 dtypes are not currently '
+                         'supported; please store the data using int64 instead')
+
+    return dtype, object_codec
+
+
 # noinspection PyTypeChecker
 def is_total_slice(item, shape):
     """Determine whether `item` specifies a complete slice of array with the