8000 BUG: Move ndarray.dump to python and make it close the file it opens · numpy/numpy@87eeb8d · GitHub
[go: up one dir, main page]
Skip to content

Commit 87eeb8d

Browse files
eric-wiesereric-wieser
committed
BUG: Move ndarray.dump to python and make it close the file it opens
As a side-effect, this makes it support kwargs.
1 parent 392866d commit 87eeb8d

File tree

2 files changed

+34
-55
lines changed

2 files changed

+34
-55
lines changed

‎numpy/core/_methods.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
from numpy.core import numerictypes as nt
1414
from numpy.core import _exceptions
1515
from numpy._globals import _NoValue
16+
from numpy.compat import pickle, os_fspath, contextlib_nullcontext
1617

1718
# save those O(100) nanoseconds!
1819
umr_maximum = um.maximum.reduce
@@ -230,3 +231,14 @@ def _ptp(a, axis=None, out=None, keepdims=False):
230231
umr_minimum(a, axis, None, None, keepdims),
231232
out
232233
)
234+
235+
def _dump(self, file, protocol=2):
236+
if hasattr(file, 'write'):
237+
ctx = contextlib_nullcontext(file)
238+
else:
239+
ctx = open(os_fspath(file), "wb")
240+
with ctx as f:
241+
pickle.dump(self, f, protocol=protocol)
242+
243+
def _dumps(self, protocol=2):
244+
return pickle.dumps(self, protocol=protocol)

‎numpy/core/src/multiarray/methods.c

Lines changed: 22 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -2098,87 +2098,54 @@ array_setstate(PyArrayObject *self, PyObject *args)
20982098
NPY_NO_EXPORT int
20992099
PyArray_Dump(PyObject *self, PyObject *file, int protocol)
21002100
{
2101-
PyObject *cpick = NULL;
2101+
static PyObject *method = NULL;
21022102
PyObject *ret;
2103-
if (protocol < 0) {
2104-
protocol = 2;
2105-
}
2106-
2107-
#if defined(NPY_PY3K)
2108-
cpick = PyImport_ImportModule("pickle");
2109-
#else
2110-
cpick = PyImport_ImportModule("cPickle");
2111-
#endif
2112-
if (cpick == NULL) {
2103+
npy_cache_import("numpy.core._methods", "_dump", &method);
2104+
if (method == NULL) {
21132105
return -1;
21142106
}
2115-
if (PyBytes_Check(file) || PyUnicode_Check(file)) {
2116-
file = npy_PyFile_OpenFile(file, "wb");
2117-
if (file == NULL) {
2118-
Py_DECREF(cpick);
2119-
return -1;
2120-
}
2107+
if (protocol < 0) {
2108+
ret = PyObject_CallFunction(method, "OO", self, file);
21212109
}
21222110
else {
2123-
Py_INCREF(file);
< 10000 /code>
2111+
ret = PyObject_CallFunction(method, "OOi", self, file, protocol);
21242112
}
2125-
ret = PyObject_CallMethod(cpick, "dump", "OOi", self, file, protocol);
2126-
Py_XDECREF(ret);
2127-
Py_DECREF(file);
2128-
Py_DECREF(cpick);
2129-
if (PyErr_Occurred()) {
2113+
if (ret == NULL) {
21302114
return -1;
21312115
}
2116+
Py_DECREF(ret);
21322117
return 0;
21332118
}
21342119

21352120
/*NUMPY_API*/
21362121
NPY_NO_EXPORT PyObject *
21372122
PyArray_Dumps(PyObject *self, int protocol)
21382123
{
2139-
PyObject *cpick = NULL;
2140-
PyObject *ret;
2124+
static PyObject *method = NULL;
2125+
npy_cache_import("numpy.core._methods", "_dumps", &method);
2126+
if (method == NULL) {
2127+
return NULL;
2128+
}
21412129
if (protocol < 0) {
2142-
protocol = 2;
2130+
return PyObject_CallFunction(method, "O", self);
21432131
}
2144-
#if defined(NPY_PY3K)
2145-
cpick = PyImport_ImportModule("pickle");
2146-
#else
2147-
cpick = PyImport_ImportModule("cPickle");
2148-
#endif
2149-
if (cpick == NULL) {
2150-
return NULL;
2132+
else {
2133+
return PyObject_CallFunction(method, "Oi", self, protocol);
21512134
}
2152-
ret = PyObject_CallMethod(cpick, "dumps", "Oi", self, protocol);
2153-
Py_DECREF(cpick);
2154-
return ret;
21552135
}
21562136

21572137

21582138
static PyObject *
2159-
array_dump(PyArrayObject *self, PyObject *args)
2139+
array_dump(PyArrayObject *self, PyObject *args, PyObject *kwds)
21602140
{
2161-
PyObject *file = NULL;
2162-
int ret;
2163-
2164-
if (!PyArg_ParseTuple(args, "O:dump", &file)) {
2165-
return NULL;
2166-
}
2167-
ret = PyArray_Dump((PyObject *)self, file, 2);
2168-
if (ret < 0) {
2169-
return NULL;
2170-
}
2171-
Py_RETURN_NONE;
2141+
NPY_FORWARD_NDARRAY_METHOD("_dump");
21722142
}
21732143

21742144

21752145
static PyObject *
2176-
array_dumps(PyArrayObject *self, PyObject *args)
2146+
array_dumps(PyArrayObject *self, PyObject *args, PyObject *kwds)
21772147
{
2178-
if (!PyArg_ParseTuple(args, "")) {
2179-
return NULL;
2180-
}
2181-
return PyArray_Dumps((PyObject *)self, 2);
2148+
NPY_FORWARD_NDARRAY_METHOD("_dumps");
21822149
}
21832150

21842151

@@ -2753,10 +2720,10 @@ NPY_NO_EXPORT PyMethodDef array_methods[] = {
27532720
METH_VARARGS, NULL},
27542721
{"dumps",
27552722
(PyCFunction) array_dumps,
2756-
METH_VARARGS, NULL},
2723+
METH_VARARGS | METH_KEYWORDS, NULL},
27572724
{"dump",
27582725
(PyCFunction) array_dump,
2759-
METH_VARARGS, NULL},
2726+
METH_VARARGS | METH_KEYWORDS, NULL},
27602727

27612728
{"__complex__",
27622729
(PyCFunction) array_complex,

0 commit comments

Comments
 (0)
0