merged

flexatone · flexatone · commit 8a2c5a978a3c · 2026-06-29T22:09:24.000-07:00
diff --git a/README.rst b/README.rst
@@ -35,6 +35,12 @@ ArrayKit requires the following:
 What is New in ArrayKit
 -------------------------
 
+1.4.0
+............
+
+Added ``write_array_to_file()``.
+
+
 1.3.1
 ............
 
diff --git a/VERSION b/VERSION
@@ -1,2 +1,2 @@
-1.3.1
+1.4.0
 
diff --git a/pyproject.toml b/pyproject.toml
@@ -19,7 +19,7 @@ authors = [
   { name = "Brandt Bucher" },
   { name = "Charles Burkland" },
 ]
-license = "MIT"
+license = { text = "MIT" }
 keywords = ["numpy", "array"]
 dependencies = [
   "numpy>=1.24.3",
diff --git a/requirements-dev-3_11.txt b/requirements-dev-3_11.txt
@@ -1,4 +1,4 @@
 numpy==2.0.2
-pytest==7.1.2
+pytest==9.0.3
 nox==2025.5.1
 hypothesis==6.10.1
diff --git a/requirements-dev-3_12.txt b/requirements-dev-3_12.txt
@@ -1,4 +1,4 @@
 numpy==2.2.6
-pytest==7.1.2
+pytest==9.0.3
 nox==2025.5.1
 hypothesis==6.10.1
diff --git a/requirements-dev-3_13.txt b/requirements-dev-3_13.txt
@@ -1,4 +1,4 @@
 numpy==2.3.1
-pytest==8.3.3
+pytest==9.0.3
 nox==2025.5.1
 hypothesis==6.131.16
diff --git a/requirements-dev-3_14.txt b/requirements-dev-3_14.txt
@@ -1,4 +1,4 @@
 numpy==2.3.3
-pytest==8.3.3
+pytest==9.0.3
 nox==2025.5.1
 hypothesis==6.131.16
diff --git a/setup.py b/setup.py
@@ -9,6 +9,13 @@
 
 def get_ext_dir(*components: tp.Iterable[str]) -> tp.Sequence[str]:
     dirs = []
+    # Check user site-packages
+    user_site = site.getusersitepackages()
+    if user_site:
+        fp = os.path.join(user_site, *components)
+        if os.path.exists(fp):
+            dirs.append(fp)
+    # Check system site-packages
     for sp in site.getsitepackages():
         fp = os.path.join(sp, *components)
         if os.path.exists(fp):
diff --git a/src/__init__.py b/src/__init__.py
@@ -24,6 +24,7 @@
 from ._arraykit import iterable_str_to_array_1d as iterable_str_to_array_1d
 from ._arraykit import split_after_count as split_after_count
 from ._arraykit import get_new_indexers_and_screen as get_new_indexers_and_screen
+from ._arraykit import write_array_to_file as write_array_to_file
 from ._arraykit import count_iteration as count_iteration
 from ._arraykit import first_true_1d as first_true_1d
 from ._arraykit import first_true_2d as first_true_2d
diff --git a/src/__init__.pyi b/src/__init__.pyi
@@ -217,6 +217,13 @@ def dtype_from_element(__value: tp.Optional[_TLabel]) -> np.dtype: ...
 def get_new_indexers_and_screen(
     indexers: np.ndarray, positions: np.ndarray
 ) -> tp.Tuple[np.ndarray, np.ndarray]: ...
+def write_array_to_file(
+    __array: np.ndarray,
+    __file: tp.IO[bytes],
+    *,
+    fortran_order: bool = False,
+    buffersize: int = 8192,
+) -> None: ...
 def first_true_1d(__array: np.ndarray, *, forward: bool) -> int: ...
 def first_true_2d(__array: np.ndarray, *, forward: bool, axis: int) -> np.ndarray: ...
 def nonzero_1d(__array: np.ndarray, /) -> np.ndarray: ...
diff --git a/src/_arraykit.c b/src/_arraykit.c
@@ -66,6 +66,10 @@ static PyMethodDef arraykit_methods[] =  {
             (PyCFunction)get_new_indexers_and_screen,
             METH_VARARGS | METH_KEYWORDS,
             NULL},
+    {"write_array_to_file",
+            (PyCFunction)write_array_to_file,
+            METH_VARARGS | METH_KEYWORDS,
+            NULL},
     {NULL},
 };
 
diff --git a/src/methods.c b/src/methods.c
@@ -9,6 +9,12 @@
 # include "numpy/halffloat.h"
 # include <string.h>
 
+# ifdef _WIN32
+#  include <io.h>
+# else
+#  include <unistd.h>
+# endif
+
 # include "methods.h"
 # include "utilities.h"
 
@@ -1400,6 +1406,248 @@ static char *array_deepcopy_kwarg_names[] = {
     NULL
 };
 
+static char *write_array_to_file_kwarg_names[] = {
+    "array",
+    "file",
+    "fortran_order",
+    "buffersize",
+    NULL
+};
+
+/* Helper function to write data directly to a file descriptor.
+ * Returns -1 on error (with Python exception set), or 0 on success.
+ */
+static int
+write_to_fd(int fd, const char *data, Py_ssize_t size)
+{
+    Py_ssize_t total_written = 0;
+    while (total_written < size) {
+        Py_ssize_t to_write = size - total_written;
+#ifdef _WIN32
+        /* On Windows, _write has a max count of INT_MAX */
+        if (to_write > INT_MAX) {
+            to_write = INT_MAX;
+        }
+        int written = _write(fd, data + total_written, (unsigned int)to_write);
+#else
+        /* On POSIX, write may have platform-specific limits */
+        ssize_t written = write(fd, data + total_written, to_write);
+#endif
+        if (written < 0) {
+            PyErr_SetFromErrno(PyExc_OSError);
+            return -1;
+        }
+        if (written == 0) {
+            /* No progress - treat as error */
+            PyErr_SetString(PyExc_OSError, "write() returned 0");
+            return -1;
+        }
+        total_written += written;
+    }
+    return 0;
+}
+
+PyObject *
+write_array_to_file(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs)
+{
+    PyObject *a;
+    PyObject *file;
+    int fortran_order = 0;
+    npy_intp buffersize = 8192;
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs,
+            "OO|pn:write_array_to_file", write_array_to_file_kwarg_names,
+            &a,
+            &file,
+            &fortran_order,
+            &buffersize)) {
+        return NULL;
+    }
+    AK_CHECK_NUMPY_ARRAY(a);
+    if (buffersize < 1) {
+        PyErr_SetString(PyExc_ValueError, "buffersize must be at least 1");
+        return NULL;
+    }
+
+    PyArrayObject *array = (PyArrayObject*)a;
+    npy_uint32 flags = NPY_ITER_EXTERNAL_LOOP | NPY_ITER_BUFFERED | NPY_ITER_ZEROSIZE_OK;
+    NPY_ORDER order = (fortran_order && !PyArray_IS_C_CONTIGUOUS(array))
+            ? NPY_FORTRANORDER
+            : NPY_CORDER;
+    PyArrayObject *operands[] = {array};
+    npy_uint32 op_flags[] = {NPY_ITER_READONLY | NPY_ITER_ALIGNED};
+    NpyIter *iter = NpyIter_AdvancedNew(
+            1,
+            operands,
+            flags,
+            order,
+            NPY_NO_CASTING,
+            op_flags,
+            NULL,
+            -1,
+            NULL,
+            NULL,
+            buffersize
+            );
+    if (iter == NULL) {
+        return NULL;
+    }
+
+    NpyIter_IterNextFunc *iternext = NpyIter_GetIterNext(iter, NULL);
+    if (iternext == NULL) {
+        NpyIter_Deallocate(iter);
+        return NULL;
+    }
+
+    char **dataptr = NpyIter_GetDataPtrArray(iter);
+    npy_intp *strideptr = NpyIter_GetInnerStrideArray(iter);
+    npy_intp *innersizeptr = NpyIter_GetInnerLoopSizePtr(iter);
+    Py_ssize_t itemsize = PyArray_ITEMSIZE(array);
+    if (itemsize == 0) {
+        NpyIter_Deallocate(iter);
+        Py_RETURN_NONE;
+    }
+
+    // Try to get file descriptor for direct writes. fall back to calling Python's write() method.
+    int fd = PyObject_AsFileDescriptor(file);
+    int use_fd = (fd >= 0);
+
+    /* If we are going to write directly to the file descriptor, we must first
+     * flush any data buffered at the Python level (e.g. an io.BufferedWriter
+     * from open(path, 'wb')). Writing to the raw fd bypasses that user-space
+     * buffer, so without a flush our bytes would land at the kernel offset,
+     * ahead of any still-buffered data, corrupting the file. Raw integer fds
+     * (e.g. from os.open) have no flush() method; that is not an error. */
+    if (use_fd) {
+        PyObject *flush_result = PyObject_CallMethod(file, "flush", NULL);
+        if (flush_result == NULL) {
+            PyErr_Clear();
+        }
+        else {
+            Py_DECREF(flush_result);
+        }
+    }
+
+    // For non-fd path, we need the write method name
+    PyObject *write_name = NULL;
+    if (!use_fd) {
+        PyErr_Clear();  /* Clear the error from PyObject_AsFileDescriptor */
+        write_name = PyUnicode_InternFromString("write");
+        if (write_name == NULL) {
+            NpyIter_Deallocate(iter);
+            return NULL;
+        }
+    }
+
+    // Allocate a buffer for packing non-contiguous data
+    char *pack_buffer = NULL;
+    Py_ssize_t pack_buffer_size = 0;
+
+    /* For the write() path, reuse a single memoryview across iterations,
+     * rebinding it to each chunk, instead of allocating a new object each time. */
+    PyObject *buffer = NULL;
+
+    do {
+        npy_intp inner_size = *innersizeptr;
+        if (inner_size == 0) {
+            continue;
+        }
+        if (inner_size > PY_SSIZE_T_MAX / itemsize) {
+            PyErr_SetString(PyExc_OverflowError, "array chunk too large");
+            goto fail;
+        }
+        Py_ssize_t chunk_size = (Py_ssize_t)inner_size * itemsize;
+
+        const char *data_to_write = *dataptr;
+
+        // If stride is not contiguous, pack into buffer
+        if (*strideptr != itemsize) {
+            if (pack_buffer_size < chunk_size) {
+                char *new_buffer = (char *)PyMem_Realloc(pack_buffer, chunk_size);
+                if (new_buffer == NULL) {
+                    PyErr_NoMemory();
+                    goto fail;
+                }
+                pack_buffer = new_buffer;
+                pack_buffer_size = chunk_size;
+            }
+            char *src = *dataptr;
+            npy_intp stride = *strideptr;
+            for (npy_intp i = 0; i < inner_size; ++i) {
+                memcpy(pack_buffer + (i * itemsize), src + (i * stride), itemsize);
+            }
+            data_to_write = pack_buffer;
+        }
+
+        // Write the data
+        if (use_fd) {
+            if (write_to_fd(fd, data_to_write, chunk_size) < 0) {
+                goto fail;
+            }
+        }
+        else {
+            /* Reuse a single memoryview, rebinding it to the current chunk. This
+             * is safe because write() consumes the bytes synchronously and does
+             * not retain the view (exports return to 0 before write() returns).
+             * Note: PyMemoryView_FromMemory requires non-const char*, but we pass
+             * PyBUF_READ flag which makes the view read-only, so the cast is safe. */
+            if (buffer == NULL) {
+                buffer = PyMemoryView_FromMemory((char *)data_to_write, chunk_size, PyBUF_READ);
+                if (buffer == NULL) {
+                    goto fail;
+                }
+            }
+            else {
+                // Rebind buf/len (and shape[0], since ndim == 1 and itemsize == 1) to this chunk
+                Py_buffer *view = PyMemoryView_GET_BUFFER(buffer);
+                view->buf = (void *)data_to_write;
+                view->len = chunk_size;
+                view->shape[0] = chunk_size;
+            }
+
+            PyObject *write_result = PyObject_CallMethodObjArgs(file, write_name, buffer, NULL);
+            if (write_result == NULL) {
+                goto fail;
+            }
+
+            // Check for partial writes
+            if (PyLong_Check(write_result)) {
+                Py_ssize_t bytes_written = PyLong_AsSsize_t(write_result);
+                if (bytes_written < 0 && PyErr_Occurred()) {
+                    Py_DECREF(write_result);
+                    goto fail;
+                }
+                if (bytes_written != chunk_size) {
+                    Py_DECREF(write_result);
+                    PyErr_Format(PyExc_OSError,
+                                 "write() returned partial write (%zd of %zd bytes)",
+                                 bytes_written, chunk_size);
+                    goto fail;
+                }
+            }
+            Py_DECREF(write_result);
+        }
+    } while(iternext(iter));
+
+    /* NpyIter can return 0 for TWO reasons: end of iteration OR error.
+     * We MUST check PyErr_Occurred() to distinguish between the two. */
+    if (PyErr_Occurred()) {
+        goto fail;
+    }
+
+    PyMem_Free(pack_buffer);
+    Py_XDECREF(buffer);
+    NpyIter_Deallocate(iter);
+    Py_XDECREF(write_name);
+    Py_RETURN_NONE;
+
+    fail:
+        PyMem_Free(pack_buffer);
+        Py_XDECREF(buffer);
+        NpyIter_Deallocate(iter);
+        Py_XDECREF(write_name);
+        return NULL;
+}
+
 PyObject *
 array_deepcopy(PyObject *m, PyObject *args, PyObject *kwargs)
 {
diff --git a/src/methods.h b/src/methods.h
@@ -78,6 +78,9 @@ isna_element(PyObject *m, PyObject *args, PyObject *kwargs);
 PyObject *
 get_new_indexers_and_screen(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs);
 
+PyObject *
+write_array_to_file(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs);
+
 // Specialized array deepcopy that stores immutable arrays in an optional memo dict that can be provided with kwargs.
 PyObject *
 array_deepcopy(PyObject *m, PyObject *args, PyObject *kwargs);
diff --git a/test/test_util.py b/test/test_util.py

Original file line number	Diff line number	Diff line change
`@@ -19,7 +19,7 @@ authors = [`
`19`	`19`	`{ name = "Brandt Bucher" },`
`20`	`20`	`{ name = "Charles Burkland" },`
`21`	`21`	`]`
`22`		`-license = "MIT"`
	`22`	`+license = { text = "MIT" }`
`23`	`23`	`keywords = ["numpy", "array"]`
`24`	`24`	`dependencies = [`
`25`	`25`	`"numpy>=1.24.3",`