Repository: jcrist/msgspec (language: python)

Issue:
Support encoding any Enum value type
### Description

Currently msgspec only supports encoding Enum values that strings or integers, throwing `msgspec.EncodeError: Only enums with int or str values are supported` when attempting to encode anything else.

As far as I can tell this affects the following encoders:
- msgpack (have not tested, only from reading source code)
- json (have tested)

It feels like any value supported by the encoder should be encoded rather than just these two specific types.

Alternatively, in the case of the JSON encoder (I haven't checked the other encoders), it should be possible to use the `enc_hook` to handle the other types. Currently this is not possible as the error is thrown before the `enc_hook` is called.

The specific use case that I'm hitting this on is in adding support for the msgspec encoder to `python-json-logger` (https://github.com/nhairs/python-json-logger/pull/12)

```python
>>> import msgspec.json
>>> e = msgspec.json.Encoder()
>>> import enum
>>> class SomeEnum(enum.Enum):
...  NULL = None
... 
>>> e.encode(SomeEnum.NULL)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
msgspec.EncodeError: Only enums with int or str values are supported
```


Discussion / hints:

Fix patch:
diff --git a/docs/source/supported-types.rst b/docs/source/supported-types.rst
index e268f26d..082e667a 100644
--- a/docs/source/supported-types.rst
+++ b/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
 
diff --git a/msgspec/_core.c b/msgspec/_core.c
index a32df857..3c1d6e11 100644
--- a/msgspec/_core.c
+++ b/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 *


Test patch:
diff --git a/tests/test_common.py b/tests/test_common.py
index 2ef8aaf2..5c0c6bf1 100644
--- a/tests/test_common.py
+++ b/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",
diff --git a/tests/test_to_builtins.py b/tests/test_to_builtins.py
index 56da76e9..e9f94cc9 100644
--- a/tests/test_to_builtins.py
+++ b/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",