Close #18264: int- and float-derived enums now converted to int or float.

4 files changed, 178 insertions, 32 deletions

diff --git a/Doc/library/json.rst b/Doc/library/json.rst
index faa7ac9293..68aef21793 100644
--- a/Doc/library/json.rst
+++ b/Doc/library/json.rst
@@ -349,23 +349,26 @@ Encoders and Decoders
 
    .. _py-to-json-table:
 
-   +-------------------+---------------+
-   | Python            | JSON          |
-   +===================+===============+
-   | dict              | object        |
-   +-------------------+---------------+
-   | list, tuple       | array         |
-   +-------------------+---------------+
-   | str               | string        |
-   +-------------------+---------------+
-   | int, float        | number        |
-   +-------------------+---------------+
-   | True              | true          |
-   +-------------------+---------------+
-   | False             | false         |
-   +-------------------+---------------+
-   | None              | null          |
-   +-------------------+---------------+
+   +----------------------------------------+---------------+
+   | Python                                 | JSON          |
+   +========================================+===============+
+   | dict                                   | object        |
+   +----------------------------------------+---------------+
+   | list, tuple                            | array         |
+   +----------------------------------------+---------------+
+   | str                                    | string        |
+   +----------------------------------------+---------------+
+   | int, float, int- & float-derived Enums | number        |
+   +----------------------------------------+---------------+
+   | True                                   | true          |
+   +----------------------------------------+---------------+
+   | False                                  | false         |
+   +----------------------------------------+---------------+
+   | None                                   | null          |
+   +----------------------------------------+---------------+
+
+   .. versionchanged:: 3.4
+      Added support for int- and float-derived Enum classes.
 
    To extend this to recognize other objects, subclass and implement a
    :meth:`default` method with another method that returns a serializable object
diff --git a/Lib/json/encoder.py b/Lib/json/encoder.py
index 39b550dbb0..05138383bc 100644
--- a/Lib/json/encoder.py
+++ b/Lib/json/encoder.py
@@ -175,6 +175,7 @@ class JSONEncoder(object):
     def encode(self, o):
         """Return a JSON string representation of a Python data structure.
 
+        >>> from json.encoder import JSONEncoder
         >>> JSONEncoder().encode({"foo": ["bar", "baz"]})
         '{"foo": ["bar", "baz"]}'
 
@@ -298,9 +299,13 @@ def _make_iterencode(markers, _default, _encoder, _indent, _floatstr,
             elif value is False:
                 yield buf + 'false'
             elif isinstance(value, int):
-                yield buf + str(value)
+                # Subclasses of int/float may override __str__, but we still
+                # want to encode them as integers/floats in JSON. One example
+                # within the standard library is IntEnum.
+                yield buf + str(int(value))
             elif isinstance(value, float):
-                yield buf + _floatstr(value)
+                # see comment above for int
+                yield buf + _floatstr(float(value))
             else:
                 yield buf
                 if isinstance(value, (list, tuple)):
@@ -346,7 +351,8 @@ def _make_iterencode(markers, _default, _encoder, _indent, _floatstr,
             # JavaScript is weakly typed for these, so it makes sense to
             # also allow them.  Many encoders seem to do something like this.
             elif isinstance(key, float):
-                key = _floatstr(key)
+                # see comment for int/float in _make_iterencode
+                key = _floatstr(float(key))
             elif key is True:
                 key = 'true'
             elif key is False:
@@ -354,7 +360,8 @@ def _make_iterencode(markers, _default, _encoder, _indent, _floatstr,
             elif key is None:
                 key = 'null'
             elif isinstance(key, int):
-                key = str(key)
+                # see comment for int/float in _make_iterencode
+                key = str(int(key))
             elif _skipkeys:
                 continue
             else:
@@ -374,9 +381,11 @@ def _make_iterencode(markers, _default, _encoder, _indent, _floatstr,
             elif value is False:
                 yield 'false'
             elif isinstance(value, int):
-                yield str(value)
+                # see comment for int/float in _make_iterencode
+                yield str(int(value))
             elif isinstance(value, float):
-                yield _floatstr(value)
+                # see comment for int/float in _make_iterencode
+                yield _floatstr(float(value))
             else:
                 if isinstance(value, (list, tuple)):
                     chunks = _iterencode_list(value, _current_indent_level)
@@ -402,9 +411,11 @@ def _make_iterencode(markers, _default, _encoder, _indent, _floatstr,
         elif o is False:
             yield 'false'
         elif isinstance(o, int):
-            yield str(o)
+            # see comment for int/float in _make_iterencode
+            yield str(int(o))
         elif isinstance(o, float):
-            yield _floatstr(o)
+            # see comment for int/float in _make_iterencode
+            yield _floatstr(float(o))
         elif isinstance(o, (list, tuple)):
             yield from _iterencode_list(o, _current_indent_level)
         elif isinstance(o, dict):
diff --git a/Lib/test/test_json/test_enum.py b/Lib/test/test_json/test_enum.py
new file mode 100644
index 0000000000..66c6480351
--- /dev/null
+++ b/Lib/test/test_json/test_enum.py
@@ -0,0 +1,81 @@
+from enum import Enum, IntEnum
+from test.test_json import PyTest, CTest
+
+SMALL = 1
+BIG = 1<<32
+HUGE = 1<<64
+REALLY_HUGE = 1<<96
+
+class BigNum(IntEnum):
+    small = SMALL
+    big = BIG
+    huge = HUGE
+    really_huge = REALLY_HUGE
+
+E = 2.718281
+PI = 3.141593
+TAU = 2 * PI
+
+class FloatNum(float, Enum):
+    e = E
+    pi = PI
+    tau = TAU
+
+class TestEnum:
+
+    def test_floats(self):
+        for enum in FloatNum:
+            self.assertEqual(self.dumps(enum), repr(enum.value))
+            self.assertEqual(float(self.dumps(enum)), enum)
+            self.assertEqual(self.loads(self.dumps(enum)), enum)
+
+    def test_ints(self):
+        for enum in BigNum:
+            self.assertEqual(self.dumps(enum), str(enum.value))
+            self.assertEqual(int(self.dumps(enum)), enum)
+            self.assertEqual(self.loads(self.dumps(enum)), enum)
+
+    def test_list(self):
+        self.assertEqual(
+                self.dumps(list(BigNum)),
+                str([SMALL, BIG, HUGE, REALLY_HUGE]),
+                )
+        self.assertEqual(self.dumps(list(FloatNum)), str([E, PI, TAU]))
+
+    def test_dict_keys(self):
+        s, b, h, r = BigNum
+        e, p, t = FloatNum
+        d = {
+            s:'tiny', b:'large', h:'larger', r:'largest',
+            e:"Euler's number", p:'pi', t:'tau',
+            }
+        nd = self.loads(self.dumps(d))
+        self.assertEqual(nd[str(SMALL)], 'tiny')
+        self.assertEqual(nd[str(BIG)], 'large')
+        self.assertEqual(nd[str(HUGE)], 'larger')
+        self.assertEqual(nd[str(REALLY_HUGE)], 'largest')
+        self.assertEqual(nd[repr(E)], "Euler's number")
+        self.assertEqual(nd[repr(PI)], 'pi')
+        self.assertEqual(nd[repr(TAU)], 'tau')
+
+    def test_dict_values(self):
+        d = dict(
+                tiny=BigNum.small,
+                large=BigNum.big,
+                larger=BigNum.huge,
+                largest=BigNum.really_huge,
+                e=FloatNum.e,
+                pi=FloatNum.pi,
+                tau=FloatNum.tau,
+                )
+        nd = self.loads(self.dumps(d))
+        self.assertEqual(nd['tiny'], SMALL)
+        self.assertEqual(nd['large'], BIG)
+        self.assertEqual(nd['larger'], HUGE)
+        self.assertEqual(nd['largest'], REALLY_HUGE)
+        self.assertEqual(nd['e'], E)
+        self.assertEqual(nd['pi'], PI)
+        self.assertEqual(nd['tau'], TAU)
+
+class TestPyEnum(TestEnum, PyTest): pass
+class TestCEnum(TestEnum, CTest): pass
diff --git a/Modules/_json.c b/Modules/_json.c
index 2e1659eb98..301bc87d27 100644
--- a/Modules/_json.c
+++ b/Modules/_json.c
@@ -116,6 +116,8 @@ raise_errmsg(char *msg, PyObject *s, Py_ssize_t end);
 static PyObject *
 encoder_encode_string(PyEncoderObject *s, PyObject *obj);
 static PyObject *
+encoder_encode_long(PyEncoderObject* s UNUSED, PyObject *obj);
+static PyObject *
 encoder_encode_float(PyEncoderObject *s, PyObject *obj);
 
 #define S_CHAR(c) (c >= ' ' && c <= '~' && c != '\\' && c != '"')
@@ -1302,13 +1304,45 @@ _encoded_const(PyObject *obj)
 }
 
 static PyObject *
+encoder_encode_long(PyEncoderObject* s UNUSED, PyObject *obj)
+{
+    /* Return the JSON representation of a PyLong and PyLong subclasses.
+       Calls int() on PyLong subclasses in case the str() was changed.
+       Added specifically to deal with IntEnum.  See Issue18264. */
+    PyObject *encoded, *longobj;
+    if (PyLong_CheckExact(obj)) {
+        encoded = PyObject_Str(obj);
+    }
+    else {
+        longobj = PyNumber_Long(obj);
+        if (longobj == NULL) {
+            PyErr_SetString(
+                    PyExc_ValueError,
+                    "Unable to coerce int subclass to int"
+                    );
+            return NULL;
+        }
+        encoded = PyObject_Str(longobj);
+        Py_DECREF(longobj);
+    }
+    return encoded;
+}
+
+
+static PyObject *
 encoder_encode_float(PyEncoderObject *s, PyObject *obj)
 {
-    /* Return the JSON representation of a PyFloat */
+    /* Return the JSON representation of a PyFloat.
+       Modified to call float() on float subclasses in case the subclass
+       changes the repr.  See Issue18264.  */
+    PyObject *encoded, *floatobj;
     double i = PyFloat_AS_DOUBLE(obj);
     if (!Py_IS_FINITE(i)) {
         if (!s->allow_nan) {
-            PyErr_SetString(PyExc_ValueError, "Out of range float values are not JSON compliant");
+            PyErr_SetString(
+                    PyExc_ValueError,
+                    "Out of range float values are not JSON compliant"
+                    );
             return NULL;
         }
         if (i > 0) {
@@ -1321,8 +1355,24 @@ encoder_encode_float(PyEncoderObject *s, PyObject *obj)
             return PyUnicode_FromString("NaN");
         }
     }
-    /* Use a better float format here? */
-    return PyObject_Repr(obj);
+    /* coerce float subclasses to float (primarily for Enum) */
+    if (PyFloat_CheckExact(obj)) {
+        /* Use a better float format here? */
+        encoded = PyObject_Repr(obj);
+    }
+    else {
+        floatobj = PyNumber_Float(obj);
+        if (floatobj == NULL) {
+            PyErr_SetString(
+                    PyExc_ValueError,
+                    "Unable to coerce float subclass to float"
+                    );
+            return NULL;
+        }
+        encoded = PyObject_Repr(floatobj);
+        Py_DECREF(floatobj);
+    }
+    return encoded;
 }
 
 static PyObject *
@@ -1366,7 +1416,7 @@ encoder_listencode_obj(PyEncoderObject *s, _PyAccu *acc,
         return _steal_accumulate(acc, encoded);
     }
     else if (PyLong_Check(obj)) {
-        PyObject *encoded = PyObject_Str(obj);
+        PyObject *encoded = encoder_encode_long(s, obj);
         if (encoded == NULL)
             return -1;
         return _steal_accumulate(acc, encoded);
@@ -1551,9 +1601,10 @@ encoder_listencode_dict(PyEncoderObject *s, _PyAccu *acc,
                 goto bail;
         }
         else if (PyLong_Check(key)) {
-            kstr = PyObject_Str(key);
-            if (kstr == NULL)
+            kstr = encoder_encode_long(s, key);
+            if (kstr == NULL) {
                 goto bail;
+            }
         }
         else if (skipkeys) {
             Py_DECREF(item);