Fix __setitem__ assignment in the case of slices with None for start values.

Author: tskisner

pshmem/shmem.py

@@ -225,19 +225,31 @@ def __setitem__(self, key, value):
             offset = list()
             if isinstance(key, slice):
                 # Just one dimension
-                offset.append(key.start)
+                if key.start is None:
+                    offset.append(0)
+                else:
+                    offset.append(key.start)
             else:
                 # Is it iterable?
                 try:
                     for k in key:
                         if isinstance(k, slice):
-                            offset.append(k.start)
+                            if k.start is None:
+                                offset.append(0)
+                            else:
+                                offset.append(k.start)
                         else:
                             # Must be an index
+                            if not isinstance(k, (int, np.integer)):
+                                msg = "[] key elements must be a slice or integer"
+                                raise ValueError(msg)
                             offset.append(k)
                 except TypeError:
                     # No- must be an index
-                    offset.append(k)
+                    if not isinstance(key, (int, np.integer)):
+                        msg = "[] key must be scalar or tuple of slices or integers"
+                        raise ValueError(msg)
+                    offset.append(key)
         self.set(value, offset=offset, fromrank=from_rank)
 
     def __iter__(self):

pshmem/test.py

@@ -102,8 +102,10 @@ def test_allocate(self):
                         try:
                             # Same as set(), but using __setitem__ with an
                             # allreduce to find which process is setting.
+                            #
                             # key as a tuple of offsets
                             shm[setoffset] = setdata
+
                             # key as a tuple slices
                             if setdata is None:
                                 shm[None] = setdata
@@ -160,6 +162,16 @@ def test_allocate(self):
                     # This should be bitwise identical, even for floats
                     nt.assert_equal(check[:, :, :], truth[:, :, :])
 
+                    # Try full array assignment with slices containing None start
+                    # values
+                    if p != self.rank:
+                        shm[None] = None
+                    else:
+                        shm[:, :, :] = local
+
+                    check[:, :, :] = shm[:, :, :]
+                    nt.assert_equal(check[:, :, :], truth[:, :, :])
+
                 # Ensure that we can reference the memory buffer from numpy without
                 # a memory copy.  The intention is that a slice of the shared memory
                 # buffer should appear as a C-contiguous ndarray whenever we slice