Support encoding any enum type whose value is supported

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

@@ -12947,22 +12947,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;
 }
@@ -13593,40 +13580,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;
 }
@@ -13795,8 +13767,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))) {
@@ -19417,17 +19387,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

@@ -604,20 +604,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",