Fix slice normalization and length determination

requirements.txt

@@ -1,2 +1,3 @@
 numpy
 fasteners
+kenjutsu>=0.4.2

requirements_dev.txt

@@ -6,6 +6,7 @@ coverage==4.3.4
 Cython==0.25.2
 fasteners==0.14.1
 flake8==3.3.0
+kenjutsu==0.4.2
 mccabe==0.6.1
 monotonic==1.2
 nose==1.3.7

requirements_rtfd.txt

@@ -5,3 +5,4 @@ numpydoc
 mock
 numpy
 cython
+kenjutsu>=0.4.2

setup.py

@@ -161,6 +161,7 @@ def run_setup(with_extensions):
         install_requires=[
             'numpy>=1.7',
             'fasteners',
+            'kenjutsu>=0.4.2',
         ],
         ext_modules=ext_modules,
         cmdclass=cmdclass,

zarr/core.py

@@ -6,6 +6,7 @@
 
 import numpy as np
 
+from kenjutsu.measure import len_slices
 
 from zarr.util import is_total_slice, normalize_array_selection, \
     get_chunk_range, human_readable_size, normalize_resize_args, \
@@ -449,8 +450,7 @@ def __getitem__(self, item):
         selection = normalize_array_selection(item, self._shape)
 
         # determine output array shape
-        out_shape = tuple(s.stop - s.start for s in selection
-                          if isinstance(s, slice))
+        out_shape = len_slices(selection)
 
         # setup output array
         out = np.empty(out_shape, dtype=self._dtype, order=self._order)
@@ -571,10 +571,8 @@ def __setitem__(self, item, value):
         selection = normalize_array_selection(item, self._shape)
 
         # check value shape
-        expected_shape = tuple(
-            s.stop - s.start for s in selection
-            if isinstance(s, slice)
-        )
+        expected_shape = len_slices(selection)
+
         if np.isscalar(value):
             pass
         elif expected_shape != value.shape:

zarr/tests/test_core.py

@@ -215,6 +215,10 @@ def test_array_2d(self):
         assert_array_equal(a[:110, :3], z[:110, :3])
         assert_array_equal(a[190:310, 3:7], z[190:310, 3:7])
         assert_array_equal(a[-110:, -3:], z[-110:, -3:])
+        assert_array_equal(a[0, ...], z[0, ...])
+        assert_array_equal(a[..., 0], z[..., 0])
+        assert_array_equal(a[10:20, ...], z[10:20, ...])
+        assert_array_equal(a[..., 3:7], z[..., 3:7])
         # single item
         assert_array_equal(a[0], z[0])
         assert_array_equal(a[-1], z[-1])

zarr/tests/test_util.py

@@ -51,7 +51,6 @@ def test_is_total_slice():
     assert_true(is_total_slice(slice(None), (100,)))
     assert_true(is_total_slice(slice(0, 100), (100,)))
     assert_false(is_total_slice(slice(0, 50), (100,)))
-    assert_false(is_total_slice(slice(0, 100, 2), (100,)))
 
     # 2D
     assert_true(is_total_slice(Ellipsis, (100, 100)))
@@ -61,7 +60,6 @@ def test_is_total_slice():
     assert_false(is_total_slice((slice(0, 100), slice(0, 50)), (100, 100)))
     assert_false(is_total_slice((slice(0, 50), slice(0, 100)), (100, 100)))
     assert_false(is_total_slice((slice(0, 50), slice(0, 50)), (100, 100)))
-    assert_false(is_total_slice((slice(0, 100, 2), slice(0, 100)), (100, 100)))
 
     with assert_raises(TypeError):
         is_total_slice('foo', (100,))
@@ -80,12 +78,12 @@ def test_normalize_axis_selection():
         normalize_axis_selection(-1000, 100)
 
     # slice
-    eq(slice(0, 100), normalize_axis_selection(slice(None), 100))
-    eq(slice(0, 100), normalize_axis_selection(slice(None, 100), 100))
-    eq(slice(0, 100), normalize_axis_selection(slice(0, None), 100))
-    eq(slice(0, 100), normalize_axis_selection(slice(0, 1000), 100))
-    eq(slice(99, 100), normalize_axis_selection(slice(-1, None), 100))
-    eq(slice(98, 99), normalize_axis_selection(slice(-2, -1), 100))
+    eq(slice(0, 100, 1), normalize_axis_selection(slice(None), 100))
+    eq(slice(0, 100, 1), normalize_axis_selection(slice(None, 100), 100))
+    eq(slice(0, 100, 1), normalize_axis_selection(slice(0, None), 100))
+    eq(slice(0, 100, 1), normalize_axis_selection(slice(0, 1000), 100))
+    eq(slice(99, 100, 1), normalize_axis_selection(slice(-1, None), 100))
+    eq(slice(98, 99, 1), normalize_axis_selection(slice(-2, -1), 100))
     with assert_raises(IndexError):
         normalize_axis_selection(slice(100, None), 100)
     with assert_raises(IndexError):
@@ -98,6 +96,9 @@ def test_normalize_axis_selection():
     with assert_raises(TypeError):
         normalize_axis_selection('foo', 100)
 
+    with assert_raises(TypeError):
+        normalize_axis_selection([0, 1], 100)
+
     with assert_raises(NotImplementedError):
         normalize_axis_selection(slice(0, 100, 2), 100)
 
@@ -108,34 +109,40 @@ def test_normalize_array_selection():
     eq((0,), normalize_array_selection(0, (100,)))
 
     # 1D, slice
-    eq((slice(0, 100),), normalize_array_selection(Ellipsis, (100,)))
-    eq((slice(0, 100),), normalize_array_selection(slice(None), (100,)))
-    eq((slice(0, 100),), normalize_array_selection(slice(None, 100), (100,)))
-    eq((slice(0, 100),), normalize_array_selection(slice(0, None), (100,)))
+    eq((slice(0, 100, 1),), normalize_array_selection(Ellipsis, (100,)))
+    eq((slice(0, 100, 1),), normalize_array_selection(slice(None), (100,)))
+    eq(
+        (slice(0, 100, 1),),
+        normalize_array_selection(slice(None, 100), (100,))
+    )
+    eq((slice(0, 100, 1),), normalize_array_selection(slice(0, None), (100,)))
 
     # 2D, single item
     eq((0, 0), normalize_array_selection((0, 0), (100, 100)))
     eq((99, 1), normalize_array_selection((-1, 1), (100, 100)))
 
     # 2D, single col/row
-    eq((0, slice(0, 100)), normalize_array_selection((0, slice(None)),
-                                                     (100, 100)))
-    eq((0, slice(0, 100)), normalize_array_selection((0,),
-                                                     (100, 100)))
-    eq((slice(0, 100), 0), normalize_array_selection((slice(None), 0),
-                                                     (100, 100)))
+    eq((0, slice(0, 100, 1)), normalize_array_selection((0, slice(None)),
+                                                        (100, 100)))
+    eq((0, slice(0, 100, 1)), normalize_array_selection((0,),
+                                                        (100, 100)))
+    eq((slice(0, 100, 1), 0), normalize_array_selection((slice(None), 0),
+                                                        (100, 100)))
 
     # 2D slice
-    eq((slice(0, 100), slice(0, 100)),
+    eq((slice(0, 100, 1), slice(0, 100, 1)),
        normalize_array_selection(Ellipsis, (100, 100)))
-    eq((slice(0, 100), slice(0, 100)),
+    eq((slice(0, 100, 1), slice(0, 100, 1)),
        normalize_array_selection(slice(None), (100, 100)))
-    eq((slice(0, 100), slice(0, 100)),
+    eq((slice(0, 100, 1), slice(0, 100, 1)),
        normalize_array_selection((slice(None), slice(None)), (100, 100)))
 
     with assert_raises(TypeError):
         normalize_array_selection('foo', (100,))
 
+    with assert_raises(TypeError):
+        normalize_array_selection(([0, 1],), (100,))
+
 
 def test_normalize_resize_args():
 

zarr/util.py

@@ -1,10 +1,13 @@
 # -*- coding: utf-8 -*-
 from __future__ import absolute_import, print_function, division
+import numbers
 import operator
 
 
 import numpy as np
 
+from kenjutsu.format import reformat_slices
+from kenjutsu.measure import len_slices
 
 from zarr.compat import integer_types, PY2, reduce
 
@@ -111,88 +114,44 @@ def is_total_slice(item, shape):
     given `shape`. Used to optimize __setitem__ operations on the Chunk
     class."""
 
-    # N.B., assume shape is normalized
-
-    if item == Ellipsis:
-        return True
-    if item == slice(None):
-        return True
-    if isinstance(item, slice):
-        item = item,
-    if isinstance(item, tuple):
-        return all(
-            (isinstance(s, slice) and
-                ((s == slice(None)) or
-                 ((s.stop - s.start == l) and (s.step in [1, None]))))
-            for s, l in zip(item, shape)
-        )
-    else:
-        raise TypeError('expected slice or tuple of slices, found %r' % item)
+    rf_item = normalize_array_selection(item, shape)
+
+    return len_slices(rf_item) == shape
 
 
 def normalize_axis_selection(item, l):
     """Convenience function to normalize a selection within a single axis
     of size `l`."""
 
-    if isinstance(item, int):
-        if item < 0:
-            # handle wraparound
-            item = l + item
-        if item > (l - 1) or item < 0:
-            raise IndexError('index out of bounds: %s' % item)
-        return item
-
-    elif isinstance(item, slice):
-        if item.step is not None and item.step != 1:
-            raise NotImplementedError('slice with step not supported')
-        start = 0 if item.start is None else item.start
-        stop = l if item.stop is None else item.stop
-        if start < 0:
-            start = l + start
-        if stop < 0:
-            stop = l + stop
-        if start < 0 or stop < 0:
-            raise IndexError('index out of bounds: %s, %s' % (start, stop))
-        if start >= l:
-            raise IndexError('index out of bounds: %s, %s' % (start, stop))
-        if stop > l:
-            stop = l
-        if stop < start:
-            raise IndexError('index out of bounds: %s, %s' % (start, stop))
-        return slice(start, stop)
+    rf_item = reformat_slices((item,), (l,))[0]
 
-    else:
-        raise TypeError('expected integer or slice, found: %r' % item)
+    if not isinstance(rf_item, (slice, numbers.Integral)):
+        raise TypeError("expected integer or slice, found: %r" % rf_item)
+
+    if isinstance(rf_item, slice) and rf_item.step != 1:
+        raise NotImplementedError("slice with step not supported")
+
+    if np.prod(len_slices((rf_item,))) == 0:
+        raise IndexError(
+            "index out of bounds: %s, %s" % (item.start, item.stop)
+        )
+
+    return rf_item
 
 
 # noinspection PyTypeChecker
 def normalize_array_selection(item, shape):
     """Convenience function to normalize a selection within an array with
     the given `shape`."""
 
-    # normalize item
-    if isinstance(item, integer_types):
-        item = (int(item),)
-    elif isinstance(item, slice):
-        item = (item,)
-    elif item == Ellipsis:
-        item = (slice(None),)
-
-    # handle tuple of indices/slices
-    if isinstance(item, tuple):
+    rf_item = reformat_slices(item, shape)
 
-        # determine start and stop indices for all axes
-        selection = tuple(normalize_axis_selection(i, l)
-                          for i, l in zip(item, shape))
+    # Only needed for constraint checks.
+    rf_item = tuple(
+        normalize_axis_selection(i, l) for i, l in zip(rf_item, shape)
+    )
 
-        # fill out selection if not completely specified
-        if len(selection) < len(shape):
-            selection += tuple(slice(0, l) for l in shape[len(selection):])
-
-        return selection
-
-    else:
-        raise TypeError('expected indices or slice, found: %r' % item)
+    return rf_item
 
 
 def get_chunk_range(selection, chunks):