Support scalar slice assignment

Cython/Compiler/ExprNodes.py

@@ -2346,6 +2346,8 @@ class IndexNode(ExprNode):
     is_memslice_copy = False
     memslice_ellipsis_noop = False
     warned_untyped_idx = False
+    # set by SingleAssignmentNode after analyse_types()
+    is_memslice_scalar_assignment = False
 
     def __init__(self, pos, index, *args, **kw):
         ExprNode.__init__(self, pos, index=index, *args, **kw)
@@ -2987,6 +2989,8 @@ class IndexNode(ExprNode):
         self.generate_subexpr_evaluation_code(code)
         if self.is_buffer_access or self.memslice_index:
             self.generate_buffer_setitem_code(rhs, code)
+        elif self.is_memslice_scalar_assignment:
+            self.generate_memoryviewslice_assign_scalar_code(rhs, code)
         elif self.memslice_slice:
             self.generate_memoryviewslice_setslice_code(rhs, code)
         elif self.is_memoryviewslice_access:
@@ -3089,6 +3093,12 @@ class IndexNode(ExprNode):
         import MemoryView
         MemoryView.copy_broadcast_memview_src_to_dst(rhs, self, code)
 
+    def generate_memoryviewslice_assign_scalar_code(self, rhs, code):
+        "memslice1[...] = 0.0 or memslice1[:] = 0.0"
+        import MemoryView
+        self.generate_evaluation_code(code)
+        MemoryView.assign_scalar(self, rhs, code)
+
     def put_nonecheck(self, code):
         code.globalstate.use_utility_code(raise_noneindex_error_utility_code)
         code.putln("if (%s) {" % code.unlikely("%s == Py_None") % self.base.result_as(PyrexTypes.py_object_type))

Cython/Compiler/MemoryView.py

@@ -447,6 +447,17 @@ def copy_broadcast_memview_src_to_dst(src, dst, code):
                                         dst.type.dtype.is_pyobject),
             dst.pos))
 
+def assign_scalar(dst, scalar, code):
+    "Assign a scalar to a slice. Both nodes must be temps."
+    verify_direct_dimensions(dst)
+    dtype = scalar.type
+    assert scalar.type.same_as(dst.type.dtype)
+
+    t = (dst.result(), dst.type.ndim,
+         dtype.declaration_code(""), scalar.result(), dtype.is_pyobject)
+    code.putln("__pyx_memoryview_slice_assign_scalar("
+                                "&%s, %d, sizeof(%s), &%s, %d);" % t)
+
 def copy_c_or_fortran_cname(memview):
     if memview.is_c_contig:
         c_or_f = 'c'

Cython/Compiler/Nodes.py

@@ -4753,6 +4753,8 @@ class SingleAssignmentNode(AssignmentNode):
             self.lhs.analyse_target_declaration(env)
 
     def analyse_types(self, env, use_temp = 0):
+        import ExprNodes
+
         self.rhs.analyse_types(env)
         self.lhs.analyse_target_types(env)
         self.lhs.gil_assignment_check(env)
@@ -4761,7 +4763,21 @@ class SingleAssignmentNode(AssignmentNode):
             self.lhs.memslice_broadcast = True
             self.rhs.memslice_broadcast = True
 
-        self.rhs = self.rhs.coerce_to(self.lhs.type, env)
+        is_index_node = isinstance(self.lhs, ExprNodes.IndexNode)
+        if (is_index_node and not self.rhs.type.is_memoryviewslice and
+            (self.lhs.memslice_slice or self.lhs.is_memslice_copy) and
+            (self.lhs.type.dtype.assignable_from(self.rhs.type) or
+             self.rhs.type.is_pyobject)):
+            # scalar slice assignment
+            self.lhs.is_memslice_scalar_assignment = True
+            #self.lhs = self.lhs.coerce_to_temp(env)
+            self.lhs.is_temp = True
+            dtype = self.lhs.type.dtype
+            use_temp = True
+        else:
+            dtype = self.lhs.type
+
+        self.rhs = self.rhs.coerce_to(dtype, env)
         if use_temp:
             self.rhs = self.rhs.coerce_to_temp(env)
 

Cython/Utility/MemoryView.pyx

@@ -101,7 +101,7 @@ cdef class array:
             if not self.data:
                 raise MemoryError("unable to allocate array data.")
 
-        self.dtype_is_object = format == 'O'
+        self.dtype_is_object = format == b'O'
 
     def __getbuffer__(self, Py_buffer *info, int flags):
         cdef int bufmode = -1
@@ -355,9 +355,9 @@ cdef class memoryview(object):
         if have_slices:
             obj = self.is_slice(value)
             if obj:
-                self.setitem_slice_assignment(index, obj)
+                self.setitem_slice_assignment(self[index], obj)
             else:
-                self.setitem_slice_assign_scalar(index, value)
+                self.setitem_slice_assign_scalar(self[index], value)
         else:
             self.setitem_indexed(index, value)
 
@@ -371,17 +371,42 @@ cdef class memoryview(object):
 
         return obj
 
-    cdef setitem_slice_assignment(self, index, src):
+    cdef setitem_slice_assignment(self, dst, src):
         cdef {{memviewslice_name}} dst_slice
         cdef {{memviewslice_name}} src_slice
 
-        dst = self[index]
         memoryview_copy_contents(get_slice_from_memview(src, &src_slice)[0],
                                  get_slice_from_memview(dst, &dst_slice)[0],
                                  src.ndim, dst.ndim, self.dtype_is_object)
 
-    cdef setitem_slice_assign_scalar(self, index, value):
-        raise ValueError("Scalar assignment currently unsupported")
+    cdef setitem_slice_assign_scalar(self, dst, value):
+        cdef int array[128]
+        cdef void *tmp = NULL
+        cdef void *item
+
+        cdef {{memviewslice_name}} tmp_slice, *dst_slice
+        dst_slice = get_slice_from_memview(dst, &tmp_slice)
+
+        if self.view.itemsize > sizeof(array):
+            tmp = malloc(self.view.itemsize)
+            if tmp == NULL:
+                raise MemoryError
+            item = tmp
+        else:
+            item = <void *> array
+
+        try:
+            self.assign_item_from_object(<char *> item, value)
+        except:
+            free(tmp)
+            raise
+
+        # It would be easy to support indirect dimensions, but it's easier
+        # to disallow :)
+        assert_direct_dimensions(self.view.suboffsets, self.view.ndim)
+        slice_assign_scalar(dst_slice, self.view.ndim, self.view.itemsize,
+                            item, self.dtype_is_object)
+        free(tmp)
 
     cdef setitem_indexed(self, index, value):
         cdef char *itemp = self.get_item_pointer(index)
@@ -597,6 +622,12 @@ cdef tuple _unellipsify(object index, int ndim):
 
     return have_slices or nslices, tuple(result)
 
+cdef assert_direct_dimensions(Py_ssize_t *suboffsets, int ndim):
+    cdef int i
+    for i in range(ndim):
+        if suboffsets[i] >= 0:
+            raise ValueError("Indirect dimensions not supported")
+
 #
 ### Slicing a memoryview
 #
@@ -1228,6 +1259,11 @@ cdef void broadcast_leading({{memviewslice_name}} *slice,
         slice.strides[i] = slice.strides[0]
         slice.suboffsets[i] = -1
 
+#
+### Take care of refcounting the objects in slices. Do this seperately from any copying,
+### to minimize acquiring the GIL
+#
+
 @cname('__pyx_memoryview_refcount_copying')
 cdef void refcount_copying({{memviewslice_name}} *dst, bint dtype_is_object,
                            int ndim, bint inc) nogil:
@@ -1260,6 +1296,33 @@ cdef void refcount_objects_in_slice(char *data, Py_ssize_t *shape,
 
         data += strides[0]
 
+#
+### Scalar to slice assignment
+#
+
+@cname('__pyx_memoryview_slice_assign_scalar')
+cdef void slice_assign_scalar({{memviewslice_name}} *dst, int ndim,
+                              size_t itemsize, void *item,
+                              bint dtype_is_object) nogil:
+    refcount_copying(dst, dtype_is_object, ndim, False)
+    _slice_assign_scalar(dst.data, dst.shape, dst.strides, ndim, itemsize, item)
+    refcount_copying(dst, dtype_is_object, ndim, True)
+
+@cname('__pyx_memoryview__slice_assign_scalar')
+cdef void _slice_assign_scalar(char *data, Py_ssize_t *shape,
+                              Py_ssize_t *strides, int ndim,
+                              size_t itemsize, void *item) nogil:
+    cdef Py_ssize_t i
+
+    for i in range(shape[0]):
+        if ndim == 1:
+            memcpy(data, item, itemsize)
+        else:
+            _slice_assign_scalar(data, shape + 1, strides + 1,
+                                ndim - 1, itemsize, item)
+
+        data += strides[0]
+
 
 ############### BufferFormatFromTypeInfo ###############
 cdef extern from *:

tests/run/memslice.pyx

@@ -1888,11 +1888,20 @@ class SingleObject(object):
     def __str__(self):
         return str(self.value)
 
+cdef _get_empty_object_slice(fill=None):
+    cdef cython.array a = cython.array((10,), sizeof(PyObject *), 'O')
+    cdef int i
+    for i in range(10):
+        (<PyObject **> a.data)[i] = <PyObject *> fill
+        Py_INCREF(fill)
+
+    assert a.dtype_is_object
+    return a
+
 @testcase
 def test_object_dtype_copying():
     """
     >>> test_object_dtype_copying()
-    True
     0
     1
     2
@@ -1908,28 +1917,17 @@ def test_object_dtype_copying():
     """
     cdef int i
 
-    none_refcount = get_refcount(None)
-
-    cdef cython.array a1 = cython.array((10,), sizeof(PyObject *), 'O')
-    cdef cython.array a2 = cython.array((10,), sizeof(PyObject *), 'O')
+    unique = object()
+    unique_refcount = get_refcount(unique)
 
-    print a1.dtype_is_object
-
-    cdef object[:] m1 = a1
-    cdef object[:] m2 = a2
+    cdef object[:] m1 = _get_empty_object_slice()
+    cdef object[:] m2 = _get_empty_object_slice()
 
     for i in range(10):
-        # Initialize to None first
-        (<PyObject **> a1.data)[i] = <PyObject *> None
-        Py_INCREF(None)
-        (<PyObject **> a2.data)[i] = <PyObject *> None
-        Py_INCREF(None)
-
-        # now set a unique object
         m1[i] = SingleObject(i)
 
     m2[...] = m1
-    del a1, a2, m1
+    del m1
 
     for i in range(10):
         print m2[i]
@@ -1940,4 +1938,48 @@ def test_object_dtype_copying():
     del m2
     print get_refcount(obj), obj
 
-    assert none_refcount == get_refcount(None)
+    assert unique_refcount == get_refcount(unique), (unique_refcount, get_refcount(unique))
+
+@testcase
+def test_scalar_slice_assignment():
+    """
+    >>> test_scalar_slice_assignment()
+    0
+    1
+    6
+    3
+    6
+    5
+    6
+    7
+    6
+    9
+    <BLANKLINE>
+    0
+    1
+    6
+    3
+    6
+    5
+    6
+    7
+    6
+    9
+    """
+    cdef int[10] a
+    cdef int[:] m = a
+
+    _test_scalar_slice_assignment(m)
+    print
+    _test_scalar_slice_assignment(<object> m)
+
+
+cdef _test_scalar_slice_assignment(slice_1d m):
+    cdef int i
+    for i in range(10):
+        m[i] = i
+
+    m[-2:0:-2] = 6
+
+    for i in range(10):
+        print m[i]