From 5099d35b1d63850388858c1d43c7e5f8f5a97c9f Mon Sep 17 00:00:00 2001 From: Sergey B Kirpichev Date: Fri, 13 Sep 2024 12:11:53 +0300 Subject: [PATCH] Add mpz.__array__() method to interact with numpy Closes #507 --- pyproject.toml | 3 +- src/gmpy2_mpz.c | 1 + src/gmpy2_mpz_misc.c | 85 ++++++++++++++++++++++++++++++++++++++++++++ src/gmpy2_mpz_misc.h | 1 + test/test_mpz.py | 15 ++++++++ 5 files changed, 104 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index d6e515a8..7b2176c0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,8 @@ content-type = 'text/x-rst' [project.optional-dependencies] docs = ['sphinx>=4', 'sphinx-rtd-theme>=1'] -tests = ['pytest', 'hypothesis', 'cython', 'mpmath', 'setuptools'] +tests = ['pytest', 'hypothesis', 'cython', 'mpmath', 'setuptools', + 'numpy; python_version>="3.10" and platform_system=="Linux"'] [project.urls] Homepage = 'https://github.com/aleaxit/gmpy' diff --git a/src/gmpy2_mpz.c b/src/gmpy2_mpz.c index c6f68278..585a26e0 100644 --- a/src/gmpy2_mpz.c +++ b/src/gmpy2_mpz.c @@ -91,6 +91,7 @@ static PyMethodDef GMPy_MPZ_methods[] = { { "__round__", (PyCFunction)GMPy_MPZ_Method_Round, METH_FASTCALL, GMPy_doc_mpz_method_round }, { "__sizeof__", GMPy_MPZ_Method_SizeOf, METH_NOARGS, GMPy_doc_mpz_method_sizeof }, { "__trunc__", GMPy_MPZ_Method_Trunc, METH_NOARGS, GMPy_doc_mpz_method_trunc }, + { "__array__", (PyCFunction)GMPy_MPZ_Method_Array, METH_FASTCALL | METH_KEYWORDS, GMPy_doc_mpz_method_array }, { "bit_clear", GMPy_MPZ_bit_clear_method, METH_O, doc_bit_clear_method }, { "bit_count", GMPy_MPZ_bit_count_method, METH_NOARGS, doc_bit_count_method }, { "bit_flip", GMPy_MPZ_bit_flip_method, METH_O, doc_bit_flip_method }, diff --git a/src/gmpy2_mpz_misc.c b/src/gmpy2_mpz_misc.c index 4beade5a..88b1e5b2 100644 --- a/src/gmpy2_mpz_misc.c +++ b/src/gmpy2_mpz_misc.c @@ -2167,3 +2167,88 @@ GMPy_MP_Method_Conjugate(PyObject *self, PyObject *args) Py_INCREF((PyObject*)self); return (PyObject*)self; } + +PyDoc_STRVAR(GMPy_doc_mpz_method_array, +"x.__array__(dtype=None, copy=None)\n"); + +static PyObject * +GMPy_MPZ_Method_Array(PyObject *self, PyObject *const *args, + Py_ssize_t nargs, PyObject *kwnames) +{ + Py_ssize_t i, nkws = 0; + int argidx[2] = {-1, -1}; + const char* kwname; + PyObject *dtype = Py_None, *copy = Py_None; + + if (nargs > 2) { + TYPE_ERROR("__array__() takes at most 2 positional arguments"); + return NULL; + } + if (nargs >= 1) { + argidx[0] = 0; + } + if (nargs == 2) { + argidx[1] = 1; + } + + if (kwnames) { + nkws = PyTuple_GET_SIZE(kwnames); + } + if (nkws > 2) { + TYPE_ERROR("__array__() takes at most 2 keyword arguments"); + return NULL; + } + for (i = 0; i < nkws; i++) { + kwname = PyUnicode_AsUTF8(PyTuple_GET_ITEM(kwnames, i)); + if (strcmp(kwname, "dtype") == 0) { + if (nargs == 0) { + argidx[0] = (int)(nargs + i); + } + else { + TYPE_ERROR("argument for __array__() given by name ('dtype') and position (1)"); + return NULL; + } + } + else if (strcmp(kwname, "copy") == 0) { + if (nargs <= 1) { + argidx[1] = (int)(nargs + i); + } + else { + TYPE_ERROR("argument for __array__() given by name ('copy') and position (2)"); + return NULL; + } + } + else { + TYPE_ERROR("got an invalid keyword argument for __array__()"); + return NULL; + } + } + + if (argidx[0] >= 0) { + dtype = args[argidx[0]]; + } + if (argidx[1] >= 0) { + copy = args[argidx[1]]; + } + + PyObject *mod = PyImport_ImportModule("numpy"); + + if (!mod) { + return NULL; + } + + PyObject *tmp_long = GMPy_PyLong_From_MPZ((MPZ_Object *)self, NULL); + + if (!tmp_long) { + Py_DECREF(mod); + return NULL; + } + + PyObject *result = PyObject_CallMethod(mod, "array", + "OO", tmp_long, dtype); + + Py_DECREF(mod); + Py_DECREF(tmp_long); + + return result; +} diff --git a/src/gmpy2_mpz_misc.h b/src/gmpy2_mpz_misc.h index 193ed471..b809f215 100644 --- a/src/gmpy2_mpz_misc.h +++ b/src/gmpy2_mpz_misc.h @@ -59,6 +59,7 @@ static PyObject * GMPy_MPZ_Method_IsProbabPrime(PyObject *self, PyObject *const static PyObject * GMPy_MPZ_Method_IsEven(PyObject *self, PyObject *other); static PyObject * GMPy_MPZ_Method_IsOdd(PyObject *self, PyObject *other); static PyObject * GMPy_MP_Method_Conjugate(PyObject *self, PyObject *args); +static PyObject * GMPy_MPZ_Method_Array(PyObject *self, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames); static PyObject * GMPy_MPZ_Function_NumDigits(PyObject *self, PyObject *const *args, Py_ssize_t nargs); static PyObject * GMPy_MPZ_Function_Iroot(PyObject *self, PyObject *const *args, Py_ssize_t nargs); diff --git a/test/test_mpz.py b/test/test_mpz.py index 0b499e69..6d150eb5 100644 --- a/test/test_mpz.py +++ b/test/test_mpz.py @@ -3,6 +3,7 @@ import pickle from fractions import Fraction +import pytest from hypothesis import assume, example, given, settings from hypothesis.strategies import booleans, integers, sampled_from from pytest import mark, raises @@ -1758,3 +1759,17 @@ def test_issue_312(): assert not is_prime(1 - 2**4423) assert all(not is_prime(-a) for a in range(8)) assert next_prime(-8) == 2 + + +def test_mpz_array(): + numpy = pytest.importorskip('numpy') + i = 5579686107214117131790972086716881 + m = gmpy2.mpz(i) + assert numpy.longdouble(m) == numpy.longdouble(i) + assert m.__array__(dtype=numpy.longdouble) == numpy.longdouble(i) + + raises(TypeError, lambda: m.__array__(1, 2, 3)) + raises(TypeError, lambda: m.__array__(dtype=None, copy=None, spam=123)) + raises(TypeError, lambda: m.__array__(int, dtype=None)) + raises(TypeError, lambda: m.__array__(int, None, copy=None)) + raises(TypeError, lambda: m.__array__(spam=123))