From 5cda7f4404958ff7d70dd9cbe9c83acd9273d566 Mon Sep 17 00:00:00 2001
From: William Woodruff <william@trailofbits.com>
Date: Thu, 21 Mar 2024 12:25:02 -0400
Subject: [PATCH] src, test: more generic enum handling (#12)

* src, test: more generic enum handling

Signed-off-by: William Woodruff <william@trailofbits.com>

* behave more like Python's json

Signed-off-by: William Woodruff <william@trailofbits.com>

* lintage

Signed-off-by: William Woodruff <william@trailofbits.com>

---------

Signed-off-by: William Woodruff <william@trailofbits.com>
---
 src/rfc8785/_impl.py | 19 ++++++++++---------
 test/test_impl.py    | 41 ++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 50 insertions(+), 10 deletions(-)

diff --git a/src/rfc8785/_impl.py b/src/rfc8785/_impl.py
index 77c5291..3137d33 100644
--- a/src/rfc8785/_impl.py
+++ b/src/rfc8785/_impl.py
@@ -11,12 +11,10 @@
 import typing
 from io import BytesIO
 
+_Scalar = typing.Union[bool, int, str, float, None]
+
 _Value = typing.Union[
-    bool,
-    int,
-    str,
-    float,
-    None,
+    _Scalar,
     typing.Sequence["_Value"],
     typing.Tuple["_Value"],
     typing.Mapping[str, "_Value"],
@@ -195,24 +193,26 @@ def dump(obj: _Value, sink: typing.IO[bytes]) -> None:
     if obj is None:
         sink.write(b"null")
     elif isinstance(obj, bool):
+        obj = bool(obj)
         if obj is True:
             sink.write(b"true")
         else:
             sink.write(b"false")
     elif isinstance(obj, int):
-        # Annoyance: int can be subclassed by types like IntEnum,
-        # which then break or change `int.__str__`. Rather than plugging
-        # these individually, we coerce back to `int`.
         obj = int(obj)
-
         if obj < _INT_MIN or obj > _INT_MAX:
             raise IntegerDomainError(obj)
         sink.write(str(obj).encode("utf-8"))
     elif isinstance(obj, str):
+        # NOTE: We don't coerce with `str(...)`` here, since that will do
+        # the wrong thing for `(str, Enum)` subtypes where `__str__` is
+        # `Enum.__str__`.
         _serialize_str(obj, sink)
     elif isinstance(obj, float):
+        obj = float(obj)
         _serialize_float(obj, sink)
     elif isinstance(obj, (list, tuple)):
+        obj = list(obj)
         if not obj:
             # Optimization for empty lists.
             sink.write(b"[]")
@@ -225,6 +225,7 @@ def dump(obj: _Value, sink: typing.IO[bytes]) -> None:
             dump(elem, sink)
         sink.write(b"]")
     elif isinstance(obj, dict):
+        obj = dict(obj)
         if not obj:
             # Optimization for empty dicts.
             sink.write(b"{}")
diff --git a/test/test_impl.py b/test/test_impl.py
index 30d820f..8dcf81e 100644
--- a/test/test_impl.py
+++ b/test/test_impl.py
@@ -6,7 +6,7 @@
 import json
 import struct
 import sys
-from enum import IntEnum
+from enum import Enum, IntEnum
 from io import BytesIO
 
 import pytest
@@ -129,6 +129,45 @@ class X(StrEnum):
     assert json.loads(raw) == ["foo", "bar", "baz"]
 
 
+def test_dumps_enum_multiple_inheritance():
+    # Manually inheriting str, Enum should also work.
+    class X(str, Enum):
+        A = "foo"
+        B = "bar"
+        C = "baz"
+
+    raw = impl.dumps([X.A, X.B, X.C])
+    assert json.loads(raw) == ["foo", "bar", "baz"]
+
+    # Same for other JSON-able enum types.
+    class Y(dict, Enum):
+        A = {"A": "foo"}
+        B = {"B": "bar"}
+        C = {"C": "baz"}
+
+    raw = impl.dumps([Y.A, Y.B, Y.C])
+    assert json.loads(raw) == [{"A": "foo"}, {"B": "bar"}, {"C": "baz"}]
+
+    class Z(int, Enum):
+        A = 1
+        B = 2
+        C = 3
+
+    raw = impl.dumps([Z.A, Z.B, Z.C])
+    assert json.loads(raw) == [1, 2, 3]
+
+
+def test_dumps_bare_enum_fails():
+    class X(Enum):
+        A = "1"
+        B = 2
+        C = 3.0
+
+    # Python's json doesn't allow this, so we don't either.
+    with pytest.raises(impl.CanonicalizationError, match="unsupported type"):
+        impl.dumps([X.A, X.B, X.C])
+
+
 def test_dumps_nonstring_key():
     with pytest.raises(impl.CanonicalizationError, match="object keys must be strings"):
         impl.dumps({1: 2, None: 3})