Add Ellipsis support to __getitem__ slicing

distarray/dist/distarray.py

@@ -168,8 +168,11 @@ def raw_getitem(arr, index):
         # to be run locally
         def get_slice(arr, index, ddpr, comm):
             from distarray.local.maps import Distribution
-            local_distribution = Distribution(comm=comm,
-                                              dim_data=ddpr[comm.Get_rank()])
+            if len(ddpr) == 0:
+                dim_data = ()
+            else:
+                dim_data = ddpr[comm.Get_rank()]
+            local_distribution = Distribution(comm=comm, dim_data=dim_data)
             result = arr.global_index.get_slice(index, local_distribution)
             return proxyize(result)
 

distarray/dist/maps.py

@@ -240,7 +240,7 @@ def from_global_dim_dict(cls, glb_dim_dict):
         self.bounds = list(zip(bounds[:-1], bounds[1:]))
 
         self.size = bounds[-1]
-        self.grid_size = len(bounds) - 1
+        self.grid_size = max(len(bounds) - 1, 1)
 
         self.comm_padding = int(glb_dim_dict.get('comm_padding', 0))
         self.boundary_padding = int(glb_dim_dict.get('boundary_padding', 0))
@@ -292,12 +292,14 @@ def slice_owners(self, idx):
         return coords if coords != [] else [0]
 
     def get_dimdicts(self):
-        grid_ranks = range(len(self.bounds))
+        bounds = self.bounds or [[0, 0]]
+        grid_ranks = range(len(bounds))
         cpadding = self.comm_padding
-        padding = [[cpadding, cpadding] for i in grid_ranks]
-        padding[0][0] = self.boundary_padding
-        padding[-1][-1] = self.boundary_padding
-        data_tuples = zip(grid_ranks, padding, self.bounds)
+        padding = [[cpadding, cpadding] for _ in grid_ranks]
+        if len(padding) > 0:
+            padding[0][0] = self.boundary_padding
+            padding[-1][-1] = self.boundary_padding
+        data_tuples = zip(grid_ranks, padding, bounds)
         # Build the result
         out = []
         for grid_rank, padding, (start, stop) in data_tuples:
@@ -709,7 +711,8 @@ def get_dim_data_per_rank(self):
             return []
         cart_dds = product(*dds)
         coord_and_dd = [zip(*cdd) for cdd in cart_dds]
-        rank_and_dd = sorted((self.rank_from_coords[c], dd) for (c, dd) in coord_and_dd)
+        rank_and_dd = sorted((self.rank_from_coords[c], dd)
+                             for (c, dd) in coord_and_dd)
         return [dd for (_, dd) in rank_and_dd]
 
     def is_compatible(self, o):

distarray/dist/tests/test_distarray.py

@@ -171,6 +171,59 @@ def test_incomplete_index_block_dist_2d(self):
         arr = self.context.fromarray(expected)
         assert_array_equal(arr[1].toarray(), expected[1])
 
+    def test_empty_slice_1d(self):
+        shape = (10,)
+        expected = numpy.random.randint(10, size=shape)
+        arr = self.context.fromarray(expected)
+        assert_array_equal(arr[100:].toarray(), expected[100:])
+
+    def test_empty_slice_2d(self):
+        shape = (10, 20)
+        expected = numpy.random.randint(10, size=shape)
+        arr = self.context.fromarray(expected)
+        assert_array_equal(arr[100:, 100:].toarray(), expected[100:, 100:])
+
+    def test_trailing_ellipsis(self):
+        shape = (2, 3, 7, 6)
+        expected = numpy.random.randint(10, size=shape)
+        arr = self.context.fromarray(expected)
+        assert_array_equal(arr[1, ...].toarray(), expected[1, ...])
+
+    def test_leading_ellipsis(self):
+        shape = (2, 3, 7, 6)
+        expected = numpy.random.randint(10, size=shape)
+        arr = self.context.fromarray(expected)
+        assert_array_equal(arr[..., 3].toarray(), expected[..., 3])
+
+    def test_multiple_ellipsis(self):
+        shape = (2, 4, 2, 4, 1, 5)
+        expected = numpy.random.randint(10, size=shape)
+        arr = self.context.fromarray(expected)
+        assert_array_equal(arr[..., 3, ..., 4].toarray(),
+                           expected[..., 3, ..., 4])
+
+    def test_vestigial_ellipsis(self):
+        shape = (1, 2, 3)
+        expected = numpy.random.randint(10, size=shape)
+        arr = self.context.fromarray(expected)
+        assert_array_equal(arr[0, :, 0, ...].toarray(),
+                           expected[0, :, 0, ...])
+
+    def test_all_ellipsis(self):
+        shape = (3, 2, 4)
+        expected = numpy.random.randint(10, size=shape)
+        arr = self.context.fromarray(expected)
+        assert_array_equal(arr[..., ..., ..., ...].toarray(),
+                           expected[..., ..., ..., ...])
+
+    @unittest.skip("Waiting on 0d-array support.")
+    def test_0d_ellipsis(self):
+        shape = ()
+        expected = numpy.random.randint(10, size=shape)
+        arr = self.context.fromarray(expected)
+        assert_array_equal(arr[...].toarray(),
+                           expected[...])
+
 
 class TestDistArrayCreationFromGlobalDimData(ContextTestCase):
 

distarray/metadata_utils.py

@@ -335,8 +335,9 @@ def positivify(index, size):
 def sanitize_indices(indices, ndim=None, shape=None):
     """Classify and sanitize `indices`.
 
-    * Wrap Integral or slice indices into tuples
-    * Classify as 'value' or 'view'
+    * Wrap naked Integral, slice, or Ellipsis indices into tuples
+    * Classify result as 'value' or 'view'
+    * Expand `Ellipsis` objects to slices
     * If the length of the tuple-ized `indices` is < ndim (and it's
       provided),  add slice(None)'s to indices until `indices` is ndim long
     * If `shape` is provided, call `positivify` on the indices
@@ -350,21 +351,42 @@ def sanitize_indices(indices, ndim=None, shape=None):
 
     Returns
     -------
-    2-tuple of (str, ndim-tuple of slices and Integral values)
+    2-tuple of (str, n-tuple of slices and Integral values)
     """
     if isinstance(indices, Integral):
         rtype, sanitized = 'value', (indices,)
-    elif isinstance(indices, slice):
+    elif isinstance(indices, slice) or indices is Ellipsis:
         rtype, sanitized = 'view', (indices,)
     elif all(isinstance(i, Integral) for i in indices):
         rtype, sanitized = 'value', indices
-    elif all(isinstance(i, Integral) or isinstance(i, slice) for i in indices):
+    elif all(isinstance(i, Integral)
+             or isinstance(i, slice)
+             or i is Ellipsis for i in indices):
         rtype, sanitized = 'view', indices
     else:
         msg = ("Index must be an Integral, a slice, or a sequence of "
                "Integrals and slices.")
         raise TypeError(msg)
 
+    if Ellipsis in sanitized:
+        if ndim is None:
+            raise RuntimeError("Can't call `sanitize_indices` on Ellipsis "
+                               "without providing `ndim`.")
+        # expand first Ellipsis
+        diff = ndim - (len(sanitized) - 1)
+        filler = (slice(None),) * diff
+        epos = sanitized.index(Ellipsis)
+        sanitized = sanitized[:epos] + filler + sanitized[epos+1:]
+
+        # remaining Ellipsis objects are just converted to slices
+        def replace_ellipsis(idx):
+            if idx is Ellipsis:
+                return slice(None)
+            else:
+                return idx
+        sanitized = tuple(replace_ellipsis(i) for i in sanitized)
+
+
     if ndim is not None:
         diff = ndim - len(sanitized)
         if diff < 0:

distarray/tests/test_metadata_utils.py

@@ -120,6 +120,25 @@ def test_too_many_indices(self):
         with self.assertRaises(IndexError):
             metadata_utils.sanitize_indices((2, 3, 4), ndim=2)
 
+    def test_trailing_ellipsis(self):
+        ndim = 5
+        tag, sanitized = metadata_utils.sanitize_indices((10, Ellipsis),
+                                                         ndim=ndim)
+        self.assertEqual(sanitized, (10,) + (slice(None),) * (ndim-1))
+
+    def test_leading_ellipsis(self):
+        ndim = 5
+        tag, sanitized = metadata_utils.sanitize_indices((Ellipsis, 10),
+                                                         ndim=ndim)
+        self.assertEqual(sanitized, (slice(None),) * (ndim-1) + (10,))
+
+    def test_multiple_ellipsis(self):
+        ndim = 6
+        tag, sanitized = metadata_utils.sanitize_indices((Ellipsis, 10,
+                                                          Ellipsis),
+                                                         ndim=ndim)
+        self.assertEqual(sanitized, (slice(None),) * 4 + (10,  slice(None)))
+
 
 class TestGridSizes(unittest.TestCase):