Merge pull request #153 from FrancescElies/chunks_it

Chunks class iterator

RELEASE_NOTES.rst

@@ -7,6 +7,10 @@ Changes from 0.9.0 to 0.10.0
 
 - Fix pickle for in-memory carrays (#193 #194 @dataisle @esc)
 
+- Implement chunks iterator, which allows the following syntax
+  ``for chunk_ in ca._chunks``, added "internal use" indicator to carray 
+  chunks attribute (#153 @FrancescElies and @esc)
+
 Changes from 0.8.1 to 0.9.0
 ===========================
 

bcolz/carray_ext.pxd

@@ -20,6 +20,7 @@ cdef class chunks(object):
     cdef object dtype, cparams, lastchunkarr
     cdef object chunk_cached
     cdef npy_intp nchunks, nchunk_cached, len
+    cdef int _iter_count
 
     cdef read_chunk(self, nchunk)
     cdef _save(self, nchunk, chunk_)

bcolz/carray_ext.pyx

@@ -730,6 +730,19 @@ cdef class chunks(object):
             scomp = schunk.read(ctbytes)
         return scomp
 
+    def __iter__(self):
+       self._iter_count = 0
+       return self
+
+    def __next__(self):
+        cdef int i
+        if self._iter_count < self.nchunks:
+            i = self._iter_count
+            self._iter_count += 1
+            return self.__getitem__(i)
+        else:
+            raise StopIteration()
+
     def __getitem__(self, nchunk):
         cdef void *decompressed
         cdef void *compressed

bcolz/tests/test_carray.py

@@ -2160,6 +2160,31 @@ def test_repr_disk_array_append(self):
         expected = self._create_expected('a')
         self.assertEqual(expected, repr(y))
 
+class chunksIterTest(MayBeDiskTest):
+    def test00(self):
+        """Testing chunk iterator"""
+        _dtype = 'int64'
+        chunklen_ = 10
+        N = 1e2 + 3
+        a = np.arange(N, dtype=_dtype)
+        b = bcolz.carray(a, dtype=_dtype, chunklen=chunklen_,
+                         rootdir=self.rootdir)
+        # print 'nchunks', b.nchunks
+        for n, chunk_ in enumerate(b.chunks):
+            # print 'chunk nr.', n, '-->', chunk_
+            assert_array_equal(chunk_[0:chunklen_],
+                               a[n * chunklen_:(n + 1) * chunklen_],
+                               "iter chunks not working correctly")
+
+        self.assertEquals(n, len(a) // chunklen_ - 1)
+
+
+class chunksIterMemoryTest(chunksIterTest, TestCase):
+    disk = False
+
+
+class chunksIterDiskTest(chunksIterTest, TestCase):
+    disk = True
 
 if __name__ == '__main__':
     unittest.main(verbosity=2)