implement array attributes

Author: alimanfoo

zarr/attrs.py

@@ -0,0 +1,116 @@
+# -*- coding: utf-8 -*-
+from __future__ import absolute_import, print_function, division
+
+
+import json
+import os
+from threading import Lock
+import fasteners
+
+
+class PersistentAttributes(object):
+
+    def __init__(self, path, mode):
+        self._path = path
+        self._mode = mode
+
+    def __getitem__(self, item):
+
+        if not os.path.exists(self._path):
+            raise KeyError(item)
+
+        with open(self._path, mode='r') as f:
+            return json.load(f)[item]
+
+    def __setitem__(self, key, value):
+
+        # handle read-only state
+        if self._mode == 'r':
+            raise ValueError('array is read-only')
+
+        # load existing data
+        if not os.path.exists(self._path):
+            d = dict()
+        else:
+            with open(self._path, mode='r') as f:
+                d = json.load(f)
+
+        # set key value
+        d[key] = value
+
+        # write modified data
+        with open(self._path, mode='w') as f:
+            json.dump(d, f)
+
+    def __delitem__(self, key):
+
+        # handle read-only state
+        if self._mode == 'r':
+            raise ValueError('array is read-only')
+
+        # load existing data
+        if not os.path.exists(self._path):
+            d = dict()
+        else:
+            with open(self._path, mode='r') as f:
+                d = json.load(f)
+
+        # delete key value
+        del d[key]
+
+        # write modified data
+        with open(self._path, mode='w') as f:
+            json.dump(d, f)
+
+    def asdict(self):
+        if not os.path.exists(self._path):
+            d = dict()
+        else:
+            with open(self._path, mode='r') as f:
+                d = json.load(f)
+        return d
+
+    def __iter__(self):
+        return iter(self.asdict())
+
+    def __len__(self):
+        return len(self.asdict())
+
+    def keys(self):
+        return self.asdict().keys()
+
+    def values(self):
+        return self.asdict().values()
+
+    def items(self):
+        return self.asdict().items()
+
+
+class SynchronizedPersistentAttributes(PersistentAttributes):
+
+    def __init__(self, path, mode):
+        super(SynchronizedPersistentAttributes, self).__init__(path, mode)
+        lock_path = self._path + '.lock'
+        self._thread_lock = Lock()
+        self._file_lock = fasteners.InterProcessLock(lock_path)
+
+    def __getitem__(self, item):
+        with self._thread_lock:
+            with self._file_lock:
+                v = super(SynchronizedPersistentAttributes, self).__getitem__(item)
+                return v
+
+    def __setitem__(self, key, value):
+        with self._thread_lock:
+            with self._file_lock:
+                super(SynchronizedPersistentAttributes, self).__setitem__(key, value)
+
+    def __delitem__(self, key):
+        with self._thread_lock:
+            with self._file_lock:
+                super(SynchronizedPersistentAttributes, self).__delitem__(key)
+
+    def asdict(self):
+        with self._thread_lock:
+            with self._file_lock:
+                return super(SynchronizedPersistentAttributes, self).asdict()

zarr/defaults.py

@@ -7,6 +7,7 @@
 shuffle = 1  # byte shuffle
 
 # for persistence
+attrpath = '__zattr__'
 metapath = '__zmeta__'
 datapath = '__zdata__'
 datasuffix = '.blosc'

zarr/ext.pxd

@@ -53,6 +53,7 @@ cdef class BaseArray:
     cdef int _clevel
     cdef int _shuffle
     cdef object _fill_value
+    cdef object _attrs
     cdef object _cdata
     # abstract methods
     cdef BaseChunk create_chunk(self, tuple cidx)

zarr/ext.pyx

@@ -21,7 +21,7 @@ from glob import glob
 import fasteners
 
 
-from zarr import util as _util
+from zarr import util as _util, attrs as _attrs
 from zarr import defaults
 
 
@@ -801,6 +801,10 @@ cdef class BaseArray:
         def __get__(self):
             return self._shuffle
 
+    property attrs:
+        def __get__(self):
+            return self._attrs
+
     # derived properties
 
     property size:
@@ -1085,6 +1089,9 @@ cdef class Array(BaseArray):
         # instantiate chunks
         self._cdata.flat = [self.create_chunk(None) for _ in self._cdata.flat]
 
+        # setup attributes container
+        self._attrs = dict()
+
         # N.B., in the current implementation, some chunks may overhang
         # the edge of the array. This is handled during the __getitem__ and
         # __setitem__ methods by setting appropriate slices on the chunks,
@@ -1182,6 +1189,9 @@ cdef class PersistentArray(BaseArray):
         # initialize chunks
         self._init_cdata()
 
+        # initialise attributes
+        self._init_attrs()
+
     def _init_cdata(self):
 
         # initialize an object array to hold pointers to chunk objects
@@ -1191,6 +1201,10 @@ cdef class PersistentArray(BaseArray):
         for cidx in itertools.product(*(range(n) for n in self._cdata_shape)):
             self._cdata[cidx] = self.create_chunk(cidx)
 
+    def _init_attrs(self):
+        attr_path = os.path.join(self._path, defaults.attrpath)
+        self._attrs = _attrs.PersistentAttributes(attr_path, mode=self._mode)
+
     def _create(self, path, shape=None, chunks=None, dtype=None,
                 cname=None, clevel=None, shuffle=None, fill_value=None):
 
@@ -1244,6 +1258,14 @@ cdef class PersistentArray(BaseArray):
             raise ValueError('dtype %r not consistent with existing %r' %
                              (dtype, self._dtype))
 
+    property path:
+        def __get__(self):
+            return self._path
+
+    property mode:
+        def __get__(self):
+            return self._mode
+
     property cbytes:
         def __get__(self):
             return sum(c.cbytes for c in self._cdata.flat)
@@ -1316,6 +1338,11 @@ cdef class SynchronizedPersistentArray(PersistentArray):
             shuffle=self._shuffle, fill_value=self._fill_value
         )
 
+    def _init_attrs(self):
+        attr_path = os.path.join(self._path, defaults.attrpath)
+        self._attrs = _attrs.SynchronizedPersistentAttributes(attr_path,
+                                                              mode=self._mode)
+
 
 ###############################################################################
 # LAZY ARRAY CLASSES                                                          #
@@ -1397,6 +1424,9 @@ cdef class LazyArray(BaseArray):
         # initialize a dictionary for chunk objects
         self._cdata = dict()
 
+        # initialise attributes
+        self._attrs = dict()
+
         # instantiate chunks - not now!
 
     property cbytes:
@@ -1547,3 +1577,8 @@ cdef class SynchronizedLazyPersistentArray(LazyPersistentArray):
             cname=self._cname, clevel=self._clevel, shuffle=self._shuffle,
             fill_value=self._fill_value
         )
+
+    def _init_attrs(self):
+        attr_path = os.path.join(self._path, defaults.attrpath)
+        self._attrs = _attrs.SynchronizedPersistentAttributes(attr_path,
+                                                              mode=self._mode)

zarr/tests/test_array.py

@@ -276,6 +276,18 @@ def test_append_2d_axis(self):
         eq((10, 10), z.chunks)
         assert_array_equal(e, z[:])
 
+    def test_attrs(self):
+        z = self.create_array(shape=(100, 100), chunks=(10, 10), dtype='i4')
+        assert hasattr(z, 'attrs')
+        with assert_raises(KeyError):
+            v = z.attrs['foo']
+        z.attrs['foo'] = 42
+        eq(42, z.attrs['foo'])
+        z.attrs['bar'] = 4.2
+        eq(4.2, z.attrs['bar'])
+        z.attrs['baz'] = 'quux'
+        eq('quux', z.attrs['baz'])
+
 
 class TestArray(TestCase, ArrayTests):
 
@@ -322,6 +334,11 @@ def _test_persistence(self, a, chunks):
         # set data
         z[:] = a
 
+        # set attributes
+        z.attrs['foo'] = 42
+        z.attrs['bar'] = 4.2
+        z.attrs['baz'] = 'quux'
+
         # open for reading
         z2 = self.create_array(path=path, mode='r')
         eq(a.shape, z2.shape)
@@ -335,12 +352,17 @@ def _test_persistence(self, a, chunks):
         assert_array_equal(z.is_initialized, z2.is_initialized)
         assert_true(np.count_nonzero(z2.is_initialized) > 0)
         assert_array_equal(a, z2[:])
+        eq(42, z2.attrs['foo'])
+        eq(4.2, z2.attrs['bar'])
+        eq('quux', z2.attrs['baz'])
 
         # check read-only
         with assert_raises(ValueError):
             z2[:] = 0
         with assert_raises(ValueError):
             z2.resize(100)
+        with assert_raises(ValueError):
+            z2.attrs['foo'] = 0
 
         # open for read/write if exists
         z3 = self.create_array(path=path, mode='r+')