Count nchunks_initialized properly

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

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)

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):