Add PyDec_GetConst() and PyDec_GetDigits().

skrah · Aug 10, 2020 · Jul 17, 2020 · Jul 18, 2020 · Jul 18, 2020 · Jul 18, 2020
commit 72f95d80087185d358981056d149dfb6a9cbce36
diff --git a/Include/decimal.h b/Include/decimal.h
@@ -81,22 +81,32 @@ typedef struct {
 
 #define PyDec_Get_INDEX 2
 #define PyDec_Get_RETURN mpd_t *
-#define PyDec_Get_ARGS (const PyObject *)
+#define PyDec_Get_ARGS (PyObject *)
 
-#define PyDec_AsUint128Triple_INDEX 3
+#define PyDec_GetConst_INDEX 3
+#define PyDec_GetConst_RETURN const mpd_t *
+#define PyDec_GetConst_ARGS (const PyObject *)
+
+#define PyDec_GetDigits_INDEX 4
+#define PyDec_GetDigits_RETURN int64_t
+#define PyDec_GetDigits_ARGS (const PyObject *)
+
+#define PyDec_AsUint128Triple_INDEX 5
 #define PyDec_AsUint128Triple_RETURN mpd_uint128_triple_t
 #define PyDec_AsUint128Triple_ARGS (const PyObject *)
 
-#define PyDec_FromUint128Triple_INDEX 4
+#define PyDec_FromUint128Triple_INDEX 6
 #define PyDec_FromUint128Triple_RETURN PyObject *
 #define PyDec_FromUint128Triple_ARGS (const mpd_uint128_triple_t *triple)
 
-#define CPYTHON_DECIMAL_MAX_API 5
+#define CPYTHON_DECIMAL_MAX_API 7
 
 #ifdef CPYTHON_DECIMAL_MODULE
 static PyDec_TypeCheck_RETURN PyDec_TypeCheck PyDec_TypeCheck_ARGS;
 static PyDec_Alloc_RETURN PyDec_Alloc PyDec_Alloc_ARGS;
 static PyDec_Get_RETURN PyDec_Get PyDec_Get_ARGS;
+static PyDec_GetConst_RETURN PyDec_GetConst PyDec_GetConst_ARGS;
+static PyDec_GetDigits_RETURN PyDec_GetDigits PyDec_GetDigits_ARGS;
 static PyDec_AsUint128Triple_RETURN PyDec_AsUint128Triple PyDec_AsUint128Triple_ARGS;
 static PyDec_FromUint128Triple_RETURN PyDec_FromUint128Triple PyDec_FromUint128Triple_ARGS;
 #else
@@ -111,6 +121,12 @@ static void **_decimal_api;
 #define PyDec_Get \
     (*(PyDec_Get_RETURN (*)PyDec_Get_ARGS) _decimal_api[PyDec_Get_INDEX])
 
+#define PyDec_GetConst \
+    (*(PyDec_GetConst_RETURN (*)PyDec_GetConst_ARGS) _decimal_api[PyDec_GetConst_INDEX])
+
+#define PyDec_GetDigits \
+    (*(PyDec_GetDigits_RETURN (*)PyDec_GetDigits_ARGS) _decimal_api[PyDec_GetDigits_INDEX])
+
 #define PyDec_AsUint128Triple \
     (*(PyDec_AsUint128Triple_RETURN (*)PyDec_AsUint128Triple_ARGS) _decimal_api[PyDec_AsUint128Triple_INDEX])
 

diff --git a/Lib/test/test_decimal.py b/Lib/test/test_decimal.py
@@ -42,6 +42,7 @@
 import inspect
 import threading
 
+from _testcapi import decimal_get_digits
 from _testcapi import decimal_as_triple
 from _testcapi import decimal_from_triple
 
@@ -4845,6 +4846,30 @@ def from_triple(triple):
             p = decimal_from_triple(ptriple)
             self.assertEqual(str(c), str(p))
 
+    def test_decimal_get_digits(self):
+        # Capsule API
+
+        d = C.Decimal("0")
+        self.assertEqual(decimal_get_digits(d), 1)
+
+        d = C.Decimal("1234567890")
+        self.assertEqual(decimal_get_digits(d), 10)
+
+        d = C.Decimal("inf")
+        self.assertEqual(decimal_get_digits(d), 0)
+
+        d = C.Decimal("NaN")
+        self.assertEqual(decimal_get_digits(d), 0)
+
+        d = C.Decimal("sNaN")
+        self.assertEqual(decimal_get_digits(d), 0)
+
+        d = C.Decimal("NaN1234567890")
+        self.assertEqual(decimal_get_digits(d), 10)
+
+        d = C.Decimal("sNaN1234567890")
+        self.assertEqual(decimal_get_digits(d), 10)
+
     def test_decimal_api_errors(self):
         # Capsule API
 
@@ -4868,6 +4893,8 @@ def test_decimal_api_errors(self):
             self.assertRaises(C.InvalidOperation, decimal_from_triple, (0, 0, 0, 2**63-1))
             self.assertRaises(C.InvalidOperation, decimal_from_triple, (0, 0, 0, -2**63))
 
+        self.assertRaises(TypeError, decimal_get_digits, "X")
+
 class CWhitebox(unittest.TestCase):
     """Whitebox testing for _decimal"""
 

diff --git a/Modules/_decimal/_decimal.c b/Modules/_decimal/_decimal.c
@@ -5577,7 +5577,7 @@ PyDec_Alloc(void)
 }
 
 static mpd_t *
-PyDec_Get(const PyObject *v)
+PyDec_Get(PyObject *v)
 {
     if (!PyDec_Check(v)) {
         PyErr_SetString(PyExc_TypeError,
@@ -5588,6 +5588,30 @@ PyDec_Get(const PyObject *v)
     return MPD(v);
 }
 
+static const mpd_t *
+PyDec_GetConst(const PyObject *v)
+{
+    if (!PyDec_Check(v)) {
+        PyErr_SetString(PyExc_TypeError,
+            "PyDec_GetConst: argument must be a Decimal");
+        return NULL;
+    }
+
+    return MPD(v);
+}
+
+static int64_t
+PyDec_GetDigits(const PyObject *v)
+{
+    if (!PyDec_Check(v)) {
+        PyErr_SetString(PyExc_TypeError,
+            "PyDec_GetConst: argument must be a Decimal");
+        return -1;
+    }
+
+    return MPD(v)->digits;
+}
+
 static mpd_uint128_triple_t
 PyDec_AsUint128Triple(const PyObject *v)
 {
@@ -5629,6 +5653,8 @@ init_api(void)
     _decimal_api[PyDec_TypeCheck_INDEX] = (void *)PyDec_TypeCheck;
     _decimal_api[PyDec_Alloc_INDEX] = (void *)PyDec_Alloc;
     _decimal_api[PyDec_Get_INDEX] = (void *)PyDec_Get;
+    _decimal_api[PyDec_GetConst_INDEX] = (void *)PyDec_GetConst;
+    _decimal_api[PyDec_GetDigits_INDEX] = (void *)PyDec_GetDigits;
     _decimal_api[PyDec_AsUint128Triple_INDEX] = (void *)PyDec_AsUint128Triple;
     _decimal_api[PyDec_FromUint128Triple_INDEX] = (void *)PyDec_FromUint128Triple;
 

diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
@@ -2708,6 +2708,28 @@ test_PyDateTime_DELTA_GET(PyObject *self, PyObject *obj)
 
 /* Test decimal API */
 static int decimal_initialized = 0;
+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)
 {
@@ -5461,6 +5483,7 @@ static PyMethodDef TestMethods[] = {
     {"datetime_check_datetime",     datetime_check_datetime,     METH_VARARGS},
     {"datetime_check_delta",     datetime_check_delta,           METH_VARARGS},
     {"datetime_check_tzinfo",     datetime_check_tzinfo,         METH_VARARGS},
+    {"decimal_get_digits",      decimal_get_digits,              METH_O},
     {"decimal_as_triple",       decimal_as_triple,               METH_O},
     {"decimal_from_triple",     decimal_from_triple,             METH_O},
     {"make_timezones_capi",     make_timezones_capi,             METH_NOARGS},