Support encoding any enum type whose value is supported

Previously we only supported encoding and decoding enums with integer or
string values. We now support encoding enums with any value that is also
a supported type. The restriction on decoding only integer or string
values remains.

Author: jcrist

docs/source/supported-types.rst

@@ -1059,9 +1059,13 @@ other than a field for one of these types.
 ------------------------------------
 
 Enum types (`enum.Enum`, `enum.IntEnum`, `enum.StrEnum`, ...) encode as their
-member *values* in all protocols. Only enums composed of all string or all
-integer values are supported. An error is raised during decoding if the value
-isn't the proper type, or doesn't match any valid member.
+member *values* in all protocols.
+
+Any enum whose *value* is a supported type may be encoded, but only enums
+composed of all string or all integer values may be decoded.
+
+An error is raised during decoding if the value isn't the proper type, or
+doesn't match any valid member.
 
 .. code-block:: python
 

msgspec/_core.c

@@ -12944,22 +12944,9 @@ mpack_encode_enum(EncoderState *self, PyObject *obj)
     if (PyUnicode_Check(obj))
         return mpack_encode_str(self, obj);
 
-    int status;
     PyObject *value = PyObject_GetAttr(obj, self->mod->str__value_);
     if (value == NULL) return -1;
-    if (PyLong_CheckExact(value)) {
-        status = mpack_encode_long(self, value);
-    }
-    else if (PyUnicode_CheckExact(value)) {
-        status = mpack_encode_str(self, value);
-    }
-    else {
-        PyErr_SetString(
-            self->mod->EncodeError,
-            "Only enums with int or str values are supported"
-        );
-        status = -1;
-    }
+    int status = mpack_encode(self, value);
     Py_DECREF(value);
     return status;
 }
@@ -13590,40 +13577,25 @@ json_encode_raw(EncoderState *self, PyObject *obj)
     return ms_write(self, raw->buf, raw->len);
 }
 
+static int json_encode_dict_key_noinline(EncoderState *, PyObject *);
+
 static int
 json_encode_enum(EncoderState *self, PyObject *obj, bool is_key)
 {
     if (PyLong_Check(obj)) {
-        if (MS_UNLIKELY(is_key)) {
-            return json_encode_long_as_str(self, obj);
-        }
-        return json_encode_long(self, obj);
+        return is_key ? json_encode_long_as_str(self, obj) : json_encode_long(self, obj);
     }
     if (PyUnicode_Check(obj)) {
         return json_encode_str(self, obj);
     }
 
-    int status;
     PyObject *value = PyObject_GetAttr(obj, self->mod->str__value_);
     if (value == NULL) return -1;
-    if (PyLong_CheckExact(value)) {
-        if (MS_UNLIKELY(is_key)) {
-            status = json_encode_long_as_str(self, value);
-        }
-        else {
-            status = json_encode_long(self, value);
-        }
-    }
-    else if (PyUnicode_CheckExact(value)) {
-        status = json_encode_str(self, value);
-    }
-    else {
-        PyErr_SetString(
-            self->mod->EncodeError,
-            "Only enums with int or str values are supported"
-        );
-        status = -1;
-    }
+
+    int status = (
+        is_key ? json_encode_dict_key_noinline(self, value) : json_encode(self, value)
+    );
+
     Py_DECREF(value);
     return status;
 }
@@ -13792,8 +13764,6 @@ json_encode_set(EncoderState *self, PyObject *obj)
     return status;
 }
 
-static int json_encode_dict_key_noinline(EncoderState *, PyObject *);
-
 static MS_INLINE int
 json_encode_dict_key(EncoderState *self, PyObject *key) {
     if (MS_LIKELY(PyUnicode_Check(key))) {
@@ -19414,17 +19384,7 @@ static PyObject * to_builtins(ToBuiltinsState *, PyObject *, bool);
 static PyObject *
 to_builtins_enum(ToBuiltinsState *self, PyObject *obj)
 {
-    PyObject *value = PyObject_GetAttr(obj, self->mod->str__value_);
-    if (value == NULL) return NULL;
-    if (PyLong_CheckExact(value) || PyUnicode_CheckExact(value)) {
-        return value;
-    }
-    Py_DECREF(value);
-    PyErr_SetString(
-        self->mod->EncodeError,
-        "Only enums with int or str values are supported"
-    );
-    return NULL;
+    return PyObject_GetAttr(obj, self->mod->str__value_);
 }
 
 static PyObject *

tests/test_common.py

@@ -601,20 +601,27 @@ class Empty(enum.Enum):
         with pytest.raises(TypeError, match="Enum types must have at least one item"):
             proto.Decoder(Empty)
 
-    def test_unsupported_type_errors(self, proto):
-        class Bad(enum.Enum):
+    def test_encode_complex(self, proto):
+        class Complex(enum.Enum):
             A = 1.5
 
-        with pytest.raises(
-            msgspec.EncodeError, match="Only enums with int or str values are supported"
-        ):
-            proto.encode(Bad.A)
+        res = proto.encode(Complex.A)
+        sol = proto.encode(1.5)
+        assert res == sol
+
+        res = proto.encode({Complex.A: 1})
+        sol = proto.encode({1.5: 1})
+        assert res == sol
+
+    def test_decode_complex_errors(self, proto):
+        class Complex(enum.Enum):
+            A = 1.5
 
         with pytest.raises(TypeError) as rec:
-            proto.Decoder(Bad)
+            proto.Decoder(Complex)
 
         assert "Enums must contain either all str or all int values" in str(rec.value)
-        assert repr(Bad) in str(rec.value)
+        assert repr(Complex) in str(rec.value)
 
     @pytest.mark.parametrize(
         "values",

tests/test_to_builtins.py

@@ -11,7 +11,7 @@
 
 import pytest
 
-from msgspec import UNSET, EncodeError, Struct, UnsetType, to_builtins, defstruct
+from msgspec import UNSET, Struct, UnsetType, to_builtins, defstruct
 
 PY310 = sys.version_info[:2] >= (3, 10)
 PY311 = sys.version_info[:2] >= (3, 11)
@@ -216,12 +216,12 @@ def test_enum(self):
         assert res == "apple"
         assert type(res) is str
 
-    def test_enum_invalid(self):
-        class Bad(enum.Enum):
+    def test_enum_complex(self):
+        class Complex(enum.Enum):
             x = (1, 2)
 
-        with pytest.raises(EncodeError, match="Only enums with int or str"):
-            to_builtins(Bad.x)
+        res = to_builtins(Complex.x)
+        assert res is Complex.x.value
 
     @pytest.mark.parametrize(
         "in_type, out_type",