/*
* This file implements assignment from a scalar to an ndarray.
*
* Written by Mark Wiebe (mwwiebe@gmail.com)
* Copyright (c) 2011 by Enthought, Inc.
*
* See LICENSE.txt for the license.
*/
#define PY_SSIZE_T_CLEAN
#include <Python.h>
#define NPY_NO_DEPRECATED_API NPY_API_VERSION
#define _MULTIARRAYMODULE
#include <numpy/ndarraytypes.h>
#include "npy_config.h"
#include "npy_pycompat.h"
#include "convert_datatype.h"
#include "methods.h"
#include "shape.h"
#include "lowlevel_strided_loops.h"
#include "array_assign.h"
/*
* Assigns the scalar value to every element of the destination raw array.
*
* Returns 0 on success, -1 on failure.
*/
NPY_NO_EXPORT int
raw_array_assign_scalar(int ndim, npy_intp *shape,
PyArray_Descr *dst_dtype, char *dst_data, npy_intp *dst_strides,
PyArray_Descr *src_dtype, char *src_data)
{
int idim;
npy_intp shape_it[NPY_MAXDIMS], dst_strides_it[NPY_MAXDIMS];
npy_intp coord[NPY_MAXDIMS];
PyArray_StridedUnaryOp *stransfer = NULL;
NpyAuxData *transferdata = NULL;
int aligned, needs_api = 0;
npy_intp src_itemsize = src_dtype->elsize;
NPY_BEGIN_THREADS_DEF;
/* Check alignment */
aligned = raw_array_is_aligned(ndim, dst_data, dst_strides,
dst_dtype->alignment);
if (!npy_is_aligned(src_data, src_dtype->alignment)) {
aligned = 0;
}
/* Use raw iteration with no heap allocation */
if (PyArray_PrepareOneRawArrayIter(
ndim, shape,
dst_data, dst_strides,
&ndim, shape_it,
&dst_data, dst_strides_it) < 0) {
return -1;
}
/* Get the function to do the casting */
if (PyArray_GetDTypeTransferFunction(aligned,
0, dst_strides_it[0],
src_dtype, dst_dtype,
0,
&stransfer, &transferdata,
&needs_api) != NPY_SUCCEED) {
return -1;
}
if (!needs_api) {
npy_intp nitems = 1, i;
for (i = 0; i < ndim; i++) {
nitems *= shape_it[i];
}
NPY_BEGIN_THREADS_THRESHOLDED(nitems);
}
NPY_RAW_ITER_START(idim, ndim, coord, shape_it) {
/* Process the innermost dimension */
stransfer(dst_data, dst_strides_it[0], src_data, 0,
shape_it[0], src_itemsize, transferdata);
} NPY_RAW_ITER_ONE_NEXT(idim, ndim, coord,
shape_it, dst_data, dst_strides_it);
NPY_END_THREADS;
NPY_AUXDATA_FREE(transferdata);
return (needs_api && PyErr_Occurred()) ? -1 : 0;
}
/*
* Assigns the scalar value to every element of the destination raw array
* where the 'wheremask' value is True.
*
* Returns 0 on success, -1 on failure.
*/
NPY_NO_EXPORT int
raw_array_wheremasked_assign_scalar(int ndim, npy_intp *shape,
PyArray_Descr *dst_dtype, char *dst_data, npy_intp *dst_strides,
PyArray_Descr *src_dtype, char *src_data,
PyArray_Descr *wheremask_dtype, char *wheremask_data,
npy_intp *wheremask_strides)
{
int idim;
npy_intp shape_it[NPY_MAXDIMS], dst_strides_it[NPY_MAXDIMS];
npy_intp wheremask_strides_it[NPY_MAXDIMS];
npy_intp coord[NPY_MAXDIMS];
PyArray_MaskedStridedUnaryOp *stransfer = NULL;
NpyAuxData *transferdata = NULL;
int aligned, needs_api = 0;
npy_intp src_itemsize = src_dtype->elsize;
NPY_BEGIN_THREADS_DEF;
/* Check alignment */
aligned = raw_array_is_aligned(ndim, dst_data, dst_strides,
dst_dtype->alignment);
if (!npy_is_aligned(src_data, src_dtype->alignment)) {
aligned = 0;
}
/* Use raw iteration with no heap allocation */
if (PyArray_PrepareTwoRawArrayIter(
ndim, shape,
dst_data, dst_strides,
wheremask_data, wheremask_strides,
&ndim, shape_it,
&dst_data, dst_strides_it,
&wheremask_data, wheremask_strides_it) < 0) {
return -1;
}
/* Get the function to do the casting */
if (PyArray_GetMaskedDTypeTransferFunction(aligned,
0, dst_strides_it[0], wheremask_strides_it[0],
src_dtype, dst_dtype, wheremask_dtype,
0,
&stransfer, &transferdata,
&needs_api) != NPY_SUCCEED) {
return -1;
}
if (!needs_api) {
npy_intp nitems = 1, i;
for (i = 0; i < ndim; i++) {
nitems *= shape_it[i];
}
NPY_BEGIN_THREADS_THRESHOLDED(nitems);
}
NPY_RAW_ITER_START(idim, ndim, coord, shape_it) {
/* Process the innermost dimension */
stransfer(dst_data, dst_strides_it[0], src_data, 0,
(npy_bool *)wheremask_data, wheremask_strides_it[0],
shape_it[0], src_itemsize, transferdata);
} NPY_RAW_ITER_TWO_NEXT(idim, ndim, coord, shape_it,
dst_data, dst_strides_it,
wheremask_data, wheremask_strides_it);
NPY_END_THREADS;
NPY_AUXDATA_FREE(transferdata);
return (needs_api && PyErr_Occurred()) ? -1 : 0;
}
/*
* Assigns a scalar value specified by 'src_dtype' and 'src_data'
* to elements of 'dst'.
*
* dst: The destination array.
* src_dtype: The data type of the source scalar.
* src_data: The memory element of the source scalar.
* wheremask: If non-NULL, a boolean mask specifying where to copy.
* casting: An exception is raised if the assignment violates this
* casting rule.
*
* This function is implemented in array_assign_scalar.c.
*
* Returns 0 on success, -1 on failure.
*/
NPY_NO_EXPORT int
PyArray_AssignRawScalar(PyArrayObject *dst,
PyArray_Descr *src_dtype, char *src_data,
PyArrayObject *wheremask,
NPY_CASTING casting)
{
int allocated_src_data = 0;
npy_longlong scalarbuffer[4];
if (PyArray_FailUnlessWriteable(dst, "assignment destination") < 0) {
return -1;
}
/* Check the casting rule */
if (!can_cast_scalar_to(src_dtype, src_data,
PyArray_DESCR(dst), casting)) {
PyObject *errmsg;
errmsg = PyUString_FromString("Cannot cast scalar from ");
PyUString_ConcatAndDel(&errmsg,
PyObject_Repr((PyObject *)src_dtype));
PyUString_ConcatAndDel(&errmsg,
PyUString_FromString(" to "));
PyUString_ConcatAndDel(&errmsg,
PyObject_Repr((PyObject *)PyArray_DESCR(dst)));
PyUString_ConcatAndDel(&errmsg,
PyUString_FromFormat(" according to the rule %s",
npy_casting_to_string(casting)));
PyErr_SetObject(PyExc_TypeError, errmsg);
Py_DECREF(errmsg);
return -1;
}
/*
* Make a copy of the src data if it's a different dtype than 'dst'
* or isn't aligned, and the destination we're copying to has
* more than one element. To avoid having to manage object lifetimes,
* we also skip this if 'dst' has an object dtype.
*/
if ((!PyArray_EquivTypes(PyArray_DESCR(dst), src_dtype) ||
!npy_is_aligned(src_data, src_dtype->alignment)) &&
PyArray_SIZE(dst) > 1 &&
!PyDataType_REFCHK(PyArray_DESCR(dst))) {
char *tmp_src_data;
/*
* Use a static buffer to store the aligned/cast version,
* or allocate some memory if more space is needed.
*/
if (sizeof(scalarbuffer) >= PyArray_DESCR(dst)->elsize) {
tmp_src_data = (char *)&scalarbuffer[0];
}
else {
tmp_src_data = PyArray_malloc(PyArray_DESCR(dst)->elsize);
if (tmp_src_data == NULL) {
PyErr_NoMemory();
goto fail;
}
allocated_src_data = 1;
}
if (PyArray_CastRawArrays(1, src_data, tmp_src_data, 0, 0,
src_dtype, PyArray_DESCR(dst), 0) != NPY_SUCCEED) {
src_data = tmp_src_data;
goto fail;
}
/* Replace src_data/src_dtype */
src_data = tmp_src_data;
src_dtype = PyArray_DESCR(dst);
}
if (wheremask == NULL) {
/* A straightforward value assignment */
/* Do the assignment with raw array iteration */
if (raw_array_assign_scalar(PyArray_NDIM(dst), PyArray_DIMS(dst),
PyArray_DESCR(dst), PyArray_DATA(dst), PyArray_STRIDES(dst),
src_dtype, src_data) < 0) {
goto fail;
}
}
else {
npy_intp wheremask_strides[NPY_MAXDIMS];
/* Broadcast the wheremask to 'dst' for raw iteration */
if (broadcast_strides(PyArray_NDIM(dst), PyArray_DIMS(dst),
PyArray_NDIM(wheremask), PyArray_DIMS(wheremask),
PyArray_STRIDES(wheremask), "where mask",
wheremask_strides) < 0) {
goto fail;
}
/* Do the masked assignment with raw array iteration */
if (raw_array_wheremasked_assign_scalar(
PyArray_NDIM(dst), PyArray_DIMS(dst),
PyArray_DESCR(dst), PyArray_DATA(dst), PyArray_STRIDES(dst),
src_dtype, src_data,
PyArray_DESCR(wheremask), PyArray_DATA(wheremask),
wheremask_strides) < 0) {
goto fail;
}
}
if (allocated_src_data) {
PyArray_free(src_data);
}
return 0;
fail:
if (allocated_src_data) {
PyArray_free(src_data);
}
return -1;
}