diff --git a/zarr/core.py b/zarr/core.py index 61ea8752d4..f7e55b1f23 100644 --- a/zarr/core.py +++ b/zarr/core.py @@ -2,6 +2,7 @@ from __future__ import absolute_import, print_function, division import operator import itertools +import re import numpy as np @@ -340,8 +341,12 @@ def nchunks(self): def nchunks_initialized(self): """The number of chunks that have been initialized with some data.""" # TODO fix bug here, need to only count chunks - return sum(1 for k in listdir(self.chunk_store, self._path) - if k not in [array_meta_key, attrs_key]) + + # key pattern for chunk keys + prog = re.compile(r'\.'.join([r'\d+'] * min(1, self.ndim))) + + # count chunk keys + return sum(1 for k in listdir(self.chunk_store, self._path) if prog.match(k)) # backwards compability initialized = nchunks_initialized diff --git a/zarr/storage.py b/zarr/storage.py index af6df563ed..4b79a174d6 100644 --- a/zarr/storage.py +++ b/zarr/storage.py @@ -641,17 +641,22 @@ def __setitem__(self, key, value): raise KeyError(key) # write to temporary file - with tempfile.NamedTemporaryFile(mode='wb', delete=False, - dir=dir_path, - prefix=file_name + '.', - suffix='.partial') as f: - f.write(value) - temp_path = f.name - - # move temporary file into place - if os.path.exists(file_path): - os.remove(file_path) - os.rename(temp_path, file_path) + temp_path = None + try: + with tempfile.NamedTemporaryFile(mode='wb', delete=False, dir=dir_path, + prefix=file_name + '.', suffix='.partial') as f: + temp_path = f.name + f.write(value) + + # move temporary file into place + if os.path.exists(file_path): + os.remove(file_path) + os.rename(temp_path, file_path) + + finally: + # clean up if temp file still exists for whatever reason + if temp_path is not None and os.path.exists(temp_path): + os.remove(temp_path) def __delitem__(self, key): path = os.path.join(self.path, key) diff --git a/zarr/tests/test_core.py b/zarr/tests/test_core.py index 89c1c31f88..5fbfd90793 100644 --- a/zarr/tests/test_core.py +++ b/zarr/tests/test_core.py @@ -635,6 +635,16 @@ def test_array_0d(self): with assert_raises(ValueError): z[...] = np.array([1, 2, 3]) + def test_nchunks_initialized(self): + + z = self.create_array(shape=100, chunks=10) + eq(0, z.nchunks_initialized) + # manually put something into the store to confuse matters + z.store['foo'] = b'bar' + eq(0, z.nchunks_initialized) + z[:] = 42 + eq(10, z.nchunks_initialized) + class TestArrayWithPath(TestArray):