Merge pull request #96 from coldeasy/int_as_string_bitcount

Adding `int_as_string_bitcount` option

1 files changed, 74 insertions, 36 deletions

diff --git a/simplejson/_speedups.c b/simplejson/_speedups.c
index f3957db..6035560 100644
--- a/simplejson/_speedups.c
+++ b/simplejson/_speedups.c
@@ -168,7 +168,8 @@ typedef struct _PyEncoderObject {
     int use_decimal;
     int namedtuple_as_object;
     int tuple_as_array;
-    int bigint_as_string;
+    PyObject *max_long_size;
+    PyObject *min_long_size;
     PyObject *item_sort_key;
     PyObject *item_sort_kw;
     int for_json;
@@ -187,6 +188,8 @@ static PyMemberDef encoder_members[] = {
     {"skipkeys", T_OBJECT, offsetof(PyEncoderObject, skipkeys_bool), READONLY, "skipkeys"},
     {"key_memo", T_OBJECT, offsetof(PyEncoderObject, key_memo), READONLY, "key_memo"},
     {"item_sort_key", T_OBJECT, offsetof(PyEncoderObject, item_sort_key), READONLY, "item_sort_key"},
+    {"max_long_size", T_OBJECT, offsetof(PyEncoderObject, max_long_size), READONLY, "max_long_size"},
+    {"min_long_size", T_OBJECT, offsetof(PyEncoderObject, min_long_size), READONLY, "min_long_size"},
     {NULL}
 };
 
@@ -197,7 +200,7 @@ JSON_ParseEncoding(PyObject *encoding);
 static PyObject *
 JSON_UnicodeFromChar(JSON_UNICHR c);
 static PyObject *
-maybe_quote_bigint(PyObject *encoded, PyObject *obj);
+maybe_quote_bigint(PyEncoderObject* s, PyObject *encoded, PyObject *obj);
 static Py_ssize_t
 ascii_char_size(JSON_UNICHR c);
 static Py_ssize_t
@@ -384,35 +387,22 @@ JSON_UnicodeFromChar(JSON_UNICHR c)
 }
 
 static PyObject *
-maybe_quote_bigint(PyObject *encoded, PyObject *obj)
+maybe_quote_bigint(PyEncoderObject* s, PyObject *encoded, PyObject *obj)
 {
-    static PyObject *big_long = NULL;
-    static PyObject *small_long = NULL;
-    if (big_long == NULL) {
-        big_long = PyLong_FromLongLong(1LL << 53);
-        if (big_long == NULL) {
-            Py_DECREF(encoded);
-            return NULL;
-        }
-    }
-    if (small_long == NULL) {
-        small_long = PyLong_FromLongLong(-1LL << 53);
-        if (small_long == NULL) {
-            Py_DECREF(encoded);
-            return NULL;
-        }
-    }
-    if (PyObject_RichCompareBool(obj, big_long, Py_GE) ||
-        PyObject_RichCompareBool(obj, small_long, Py_LE)) {
+    if (s->max_long_size != Py_None && s->min_long_size != Py_None) {
+        if (PyObject_RichCompareBool(obj, s->max_long_size, Py_GE) ||
+            PyObject_RichCompareBool(obj, s->min_long_size, Py_LE)) {
 #if PY_MAJOR_VERSION >= 3
-        PyObject* quoted = PyUnicode_FromFormat("\"%U\"", encoded);
+            PyObject* quoted = PyUnicode_FromFormat("\"%U\"", encoded);
 #else
-        PyObject* quoted = PyString_FromFormat("\"%s\"",
-                                               PyString_AsString(encoded));
+            PyObject* quoted = PyString_FromFormat("\"%s\"",
+                                                   PyString_AsString(encoded));
 #endif
-        Py_DECREF(encoded);
-        encoded = quoted;
+            Py_DECREF(encoded);
+            encoded = quoted;
+        }
     }
+
     return encoded;
 }
 
@@ -2568,6 +2558,8 @@ encoder_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
         s->item_sort_key = NULL;
         s->item_sort_kw = NULL;
         s->Decimal = NULL;
+        s->max_long_size = NULL;
+        s->min_long_size = NULL;
     }
     return (PyObject *)s;
 }
@@ -2576,13 +2568,33 @@ static int
 encoder_init(PyObject *self, PyObject *args, PyObject *kwds)
 {
     /* initialize Encoder object */
-    static char *kwlist[] = {"markers", "default", "encoder", "indent", "key_separator", "item_separator", "sort_keys", "skipkeys", "allow_nan", "key_memo", "use_decimal", "namedtuple_as_object", "tuple_as_array", "bigint_as_string", "item_sort_key", "encoding", "for_json", "ignore_nan", "Decimal", NULL};
+    static char *kwlist[] = {
+        "markers",
+        "default",
+        "encoder",
+        "indent",
+        "key_separator",
+        "item_separator",
+        "sort_keys",
+        "skipkeys",
+        "allow_nan",
+        "key_memo",
+        "use_decimal",
+        "namedtuple_as_object",
+        "tuple_as_array",
+        "int_as_string_bitcount",
+        "item_sort_key",
+        "encoding",
+        "for_json",
+        "ignore_nan",
+        "Decimal",
+        NULL};
 
     PyEncoderObject *s;
     PyObject *markers, *defaultfn, *encoder, *indent, *key_separator;
     PyObject *item_separator, *sort_keys, *skipkeys, *allow_nan, *key_memo;
     PyObject *use_decimal, *namedtuple_as_object, *tuple_as_array;
-    PyObject *bigint_as_string, *item_sort_key, *encoding, *for_json;
+    PyObject *int_as_string_bitcount, *item_sort_key, *encoding, *for_json;
     PyObject *ignore_nan, *Decimal;
 
     assert(PyEncoder_Check(self));
@@ -2591,8 +2603,9 @@ encoder_init(PyObject *self, PyObject *args, PyObject *kwds)
     if (!PyArg_ParseTupleAndKeywords(args, kwds, "OOOOOOOOOOOOOOOOOOO:make_encoder", kwlist,
         &markers, &defaultfn, &encoder, &indent, &key_separator, &item_separator,
         &sort_keys, &skipkeys, &allow_nan, &key_memo, &use_decimal,
-        &namedtuple_as_object, &tuple_as_array, &bigint_as_string,
-        &item_sort_key, &encoding, &for_json, &ignore_nan, &Decimal))
+        &namedtuple_as_object, &tuple_as_array,
+        &int_as_string_bitcount, &item_sort_key, &encoding, &for_json,
+        &ignore_nan, &Decimal))
         return -1;
 
     s->markers = markers;
@@ -2614,7 +2627,30 @@ encoder_init(PyObject *self, PyObject *args, PyObject *kwds)
     s->use_decimal = PyObject_IsTrue(use_decimal);
     s->namedtuple_as_object = PyObject_IsTrue(namedtuple_as_object);
     s->tuple_as_array = PyObject_IsTrue(tuple_as_array);
-    s->bigint_as_string = PyObject_IsTrue(bigint_as_string);
+
+    s->max_long_size = Py_None;
+    Py_INCREF(Py_None);
+    s->min_long_size = Py_None;
+    Py_INCREF(Py_None);
+    if (PyInt_Check(int_as_string_bitcount) || PyLong_Check(int_as_string_bitcount)) {
+        static const unsigned int long_long_bitsize = SIZEOF_LONG_LONG * 8;
+        int int_as_string_bitcount_val = PyLong_AsLong(int_as_string_bitcount);
+        if (int_as_string_bitcount_val > 0 && int_as_string_bitcount_val < long_long_bitsize) {
+            Py_DECREF(Py_None);
+            Py_DECREF(Py_None);
+            s->max_long_size = PyLong_FromUnsignedLongLong(1LLU << int_as_string_bitcount_val);
+            s->min_long_size = PyLong_FromLongLong(-1LL << int_as_string_bitcount_val);
+            if (s->min_long_size == NULL || s->max_long_size == NULL) {
+                return -1;
+            }
+        }
+        else {
+            PyErr_Format(PyExc_TypeError,
+                         "int_as_string_bitcount (%d) must be greater than 0 and less than the number of bits of a `long long` type (%u bits)",
+                         int_as_string_bitcount_val, long_long_bitsize);
+        }
+    }
+
     if (item_sort_key != Py_None) {
         if (!PyCallable_Check(item_sort_key)) {
             PyErr_SetString(PyExc_TypeError, "item_sort_key must be None or callable");
@@ -2803,11 +2839,9 @@ encoder_listencode_obj(PyEncoderObject *s, JSON_Accu *rval, PyObject *obj, Py_ss
         else if (PyInt_Check(obj) || PyLong_Check(obj)) {
             PyObject *encoded = PyObject_Str(obj);
             if (encoded != NULL) {
-                if (s->bigint_as_string) {
-                    encoded = maybe_quote_bigint(encoded, obj);
-                    if (encoded == NULL)
-                        break;
-                }
+                encoded = maybe_quote_bigint(s, encoded, obj);
+                if (encoded == NULL)
+                    break;
                 rv = _steal_accumulate(rval, encoded);
             }
         }
@@ -3159,6 +3193,8 @@ encoder_traverse(PyObject *self, visitproc visit, void *arg)
     Py_VISIT(s->sort_keys);
     Py_VISIT(s->item_sort_kw);
     Py_VISIT(s->item_sort_key);
+    Py_VISIT(s->max_long_size);
+    Py_VISIT(s->min_long_size);
     Py_VISIT(s->Decimal);
     return 0;
 }
@@ -3182,6 +3218,8 @@ encoder_clear(PyObject *self)
     Py_CLEAR(s->sort_keys);
     Py_CLEAR(s->item_sort_kw);
     Py_CLEAR(s->item_sort_key);
+    Py_CLEAR(s->max_long_size);
+    Py_CLEAR(s->min_long_size);
     Py_CLEAR(s->Decimal);
     return 0;
 }