bpo-41324 Add a minimal decimal capsule API (#21519)

1 files changed, 253 insertions, 0 deletions

diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
index fca94a83a5d..593034ef65e 100644
--- a/Modules/_testcapimodule.c
+++ b/Modules/_testcapimodule.c
@@ -19,6 +19,7 @@
 
 #include "Python.h"
 #include "datetime.h"
+#include "pydecimal.h"
 #include "marshal.h"
 #include "structmember.h"         // PyMemberDef
 #include <float.h>
@@ -2705,6 +2706,252 @@ test_PyDateTime_DELTA_GET(PyObject *self, PyObject *obj)
     return Py_BuildValue("(lll)", days, seconds, microseconds);
 }
 
+/* Test decimal API */
+static int decimal_initialized = 0;
+static PyObject *
+decimal_is_special(PyObject *module, PyObject *dec)
+{
+    int is_special;
+
+    (void)module;
+    if (!decimal_initialized) {
+       if (import_decimal() < 0) {
+            return NULL;
+       }
+
+       decimal_initialized = 1;
+    }
+
+    is_special = PyDec_IsSpecial(dec);
+    if (is_special < 0) {
+        return NULL;
+    }
+
+    return PyBool_FromLong(is_special);
+}
+
+static PyObject *
+decimal_is_nan(PyObject *module, PyObject *dec)
+{
+    int is_nan;
+
+    (void)module;
+    if (!decimal_initialized) {
+       if (import_decimal() < 0) {
+            return NULL;
+       }
+
+       decimal_initialized = 1;
+    }
+
+    is_nan = PyDec_IsNaN(dec);
+    if (is_nan < 0) {
+        return NULL;
+    }
+
+    return PyBool_FromLong(is_nan);
+}
+
+static PyObject *
+decimal_is_infinite(PyObject *module, PyObject *dec)
+{
+    int is_infinite;
+
+    (void)module;
+    if (!decimal_initialized) {
+       if (import_decimal() < 0) {
+            return NULL;
+       }
+
+       decimal_initialized = 1;
+    }
+
+    is_infinite = PyDec_IsInfinite(dec);
+    if (is_infinite < 0) {
+        return NULL;
+    }
+
+    return PyBool_FromLong(is_infinite);
+}
+
+static PyObject *
+decimal_get_digits(PyObject *module, PyObject *dec)
+{
+    int64_t digits;
+
+    (void)module;
+    if (!decimal_initialized) {
+       if (import_decimal() < 0) {
+            return NULL;
+       }
+
+       decimal_initialized = 1;
+    }
+
+    digits = PyDec_GetDigits(dec);
+    if (digits < 0) {
+        return NULL;
+    }
+
+    return PyLong_FromLongLong(digits);
+}
+
+static PyObject *
+decimal_as_triple(PyObject *module, PyObject *dec)
+{
+    PyObject *tuple = NULL;
+    PyObject *sign, *hi, *lo;
+    mpd_uint128_triple_t triple;
+
+    (void)module;
+    if (!decimal_initialized) {
+       if (import_decimal() < 0) {
+            return NULL;
+       }
+
+       decimal_initialized = 1;
+    }
+
+    triple = PyDec_AsUint128Triple(dec);
+    if (triple.tag == MPD_TRIPLE_ERROR && PyErr_Occurred()) {
+        return NULL;
+    }
+
+    sign = PyLong_FromUnsignedLong(triple.sign);
+    if (sign == NULL) {
+        return NULL;
+    }
+
+    hi = PyLong_FromUnsignedLongLong(triple.hi);
+    if (hi == NULL) {
+        Py_DECREF(sign);
+        return NULL;
+    }
+
+    lo = PyLong_FromUnsignedLongLong(triple.lo);
+    if (lo == NULL) {
+        Py_DECREF(hi);
+        Py_DECREF(sign);
+        return NULL;
+    }
+
+    switch (triple.tag) {
+    case MPD_TRIPLE_QNAN:
+        assert(triple.exp == 0);
+        tuple = Py_BuildValue("(OOOs)", sign, hi, lo, "n");
+        break;
+
+    case MPD_TRIPLE_SNAN:
+        assert(triple.exp == 0);
+        tuple = Py_BuildValue("(OOOs)", sign, hi, lo, "N");
+        break;
+
+    case MPD_TRIPLE_INF:
+        assert(triple.hi == 0);
+        assert(triple.lo == 0);
+        assert(triple.exp == 0);
+        tuple = Py_BuildValue("(OOOs)", sign, hi, lo, "F");
+        break;
+
+    case MPD_TRIPLE_NORMAL:
+        tuple = Py_BuildValue("(OOOL)", sign, hi, lo, triple.exp);
+        break;
+
+    case MPD_TRIPLE_ERROR:
+        PyErr_SetString(PyExc_ValueError,
+            "value out of bounds for a uint128 triple");
+        break;
+
+    default:
+        PyErr_SetString(PyExc_RuntimeError,
+            "decimal_as_triple: internal error: unexpected tag");
+        break;
+    }
+
+    Py_DECREF(lo);
+    Py_DECREF(hi);
+    Py_DECREF(sign);
+
+    return tuple;
+}
+
+static PyObject *
+decimal_from_triple(PyObject *module, PyObject *tuple)
+{
+    mpd_uint128_triple_t triple = { MPD_TRIPLE_ERROR, 0, 0, 0, 0 };
+    PyObject *exp;
+    unsigned long sign;
+
+    (void)module;
+    if (!decimal_initialized) {
+       if (import_decimal() < 0) {
+            return NULL;
+       }
+
+       decimal_initialized = 1;
+    }
+
+    if (!PyTuple_Check(tuple)) {
+        PyErr_SetString(PyExc_TypeError, "argument must be a tuple");
+        return NULL;
+    }
+
+    if (PyTuple_GET_SIZE(tuple) != 4) {
+        PyErr_SetString(PyExc_ValueError, "tuple size must be 4");
+        return NULL;
+    }
+
+    sign = PyLong_AsUnsignedLong(PyTuple_GET_ITEM(tuple, 0));
+    if (sign == (unsigned long)-1 && PyErr_Occurred()) {
+        return NULL;
+    }
+    if (sign > UINT8_MAX) {
+        PyErr_SetString(PyExc_ValueError, "sign must be 0 or 1");
+        return NULL;
+    }
+    triple.sign = (uint8_t)sign;
+
+    triple.hi = PyLong_AsUnsignedLongLong(PyTuple_GET_ITEM(tuple, 1));
+    if (triple.hi == (unsigned long long)-1 && PyErr_Occurred()) {
+        return NULL;
+    }
+
+    triple.lo = PyLong_AsUnsignedLongLong(PyTuple_GET_ITEM(tuple, 2));
+    if (triple.lo == (unsigned long long)-1 && PyErr_Occurred()) {
+        return NULL;
+    }
+
+    exp = PyTuple_GET_ITEM(tuple, 3);
+    if (PyLong_Check(exp)) {
+        triple.tag = MPD_TRIPLE_NORMAL;
+        triple.exp = PyLong_AsLongLong(exp);
+        if (triple.exp == -1 && PyErr_Occurred()) {
+            return NULL;
+        }
+    }
+    else if (PyUnicode_Check(exp)) {
+        if (PyUnicode_CompareWithASCIIString(exp, "F") == 0) {
+            triple.tag = MPD_TRIPLE_INF;
+        }
+        else if (PyUnicode_CompareWithASCIIString(exp, "n") == 0) {
+            triple.tag = MPD_TRIPLE_QNAN;
+        }
+        else if (PyUnicode_CompareWithASCIIString(exp, "N") == 0) {
+            triple.tag = MPD_TRIPLE_SNAN;
+        }
+        else {
+            PyErr_SetString(PyExc_ValueError, "not a valid exponent");
+            return NULL;
+        }
+    }
+    else {
+        PyErr_SetString(PyExc_TypeError, "exponent must be int or string");
+        return NULL;
+    }
+
+    return PyDec_FromUint128Triple(&triple);
+}
+
 /* test_thread_state spawns a thread of its own, and that thread releases
  * `thread_done` when it's finished.  The driver code has to know when the
  * thread finishes, because the thread uses a PyObject (the callable) that
@@ -5314,6 +5561,12 @@ static PyMethodDef TestMethods[] = {
     {"PyDateTime_DATE_GET",        test_PyDateTime_DATE_GET,      METH_O},
     {"PyDateTime_TIME_GET",        test_PyDateTime_TIME_GET,      METH_O},
     {"PyDateTime_DELTA_GET",       test_PyDateTime_DELTA_GET,     METH_O},
+    {"decimal_is_special",      decimal_is_special,              METH_O},
+    {"decimal_is_nan",          decimal_is_nan,                  METH_O},
+    {"decimal_is_infinite",     decimal_is_infinite,             METH_O},
+    {"decimal_get_digits",      decimal_get_digits,              METH_O},
+    {"decimal_as_triple",       decimal_as_triple,               METH_O},
+    {"decimal_from_triple",     decimal_from_triple,             METH_O},
     {"test_list_api",           test_list_api,                   METH_NOARGS},
     {"test_dict_iteration",     test_dict_iteration,             METH_NOARGS},
     {"dict_getitem_knownhash",  dict_getitem_knownhash,          METH_VARARGS},