Skip to content

Commit a1ebbcd

Browse files
binggao1230binggao1230
authored
Fix json.encode of a dict keyed by a plain str Enum (#1118)
### Problem Encoding a dict whose keys are members of a plain `enum.Enum` with `str` values raises, even though the same key type round-trips through every other path: ```python import msgspec, enum class Fruit(enum.Enum): APPLE = "apple" BANANA = "banana" msgspec.json.decode(b'{"apple": 1}', type=dict[Fruit, int]) # ok msgspec.msgpack.encode({Fruit.APPLE: 1}) # ok msgspec.to_builtins({Fruit.APPLE: 1}) # ok -> {'apple': 1} msgspec.json.encode({Fruit.APPLE: 1}) # TypeError: Only dicts with str-like or number-like keys are supported ``` `dict[Fruit, int]` is an accepted *decode* type and both `msgpack.encode` and `to_builtins` handle the key, so `json.encode` rejecting it is a cross-codec inconsistency. ### Cause `json_encode_enum` routes a key's underlying value through `json_encode_dict_key_noinline`, which only handles non-`str` keys (`str` keys are fast-pathed in `json_encode_dict_key`). A plain `str`-valued enum therefore hits the "unsupported key" error. ### Fix When encoding an enum as a dict key, write a `str` value straight to the key with `json_encode_str`; non-`str` values keep going through `json_encode_dict_key_noinline` as before. ### Tests `test_encode_dict_plain_enum_str_key` covers a plain str `Enum` used as a dict key (encode output and encode/decode round-trip), and `FruitStr` is added to the existing mixed-key-type dict test. `pytest tests/unit/test_json.py` → 1786 passed.
1 parent 9fdf482 commit a1ebbcd

2 files changed

Lines changed: 54 additions & 3 deletions

File tree

‎src/msgspec/_core.c‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14153,9 +14153,20 @@ json_encode_enum(EncoderState *self, PyObject *obj, bool is_key)
1415314153
PyObject *value = PyObject_GetAttr(obj, self->mod->str__value_);
1415414154
if (value == NULL) return -1;
1415514155

14156-
int status = (
14157-
is_key ? json_encode_dict_key_noinline(self, value) : json_encode(self, value)
14158-
);
14156+
int status;
14157+
if (is_key) {
14158+
/* A str value must be written as a JSON object key directly.
14159+
* json_encode_dict_key_noinline only handles non-str keys (str keys
14160+
* are fast-pathed in json_encode_dict_key), so recursing into it with
14161+
* a str value would wrongly raise "unsupported key" - e.g. for a plain
14162+
* Enum whose members have str values when used as a dict key. */
14163+
status = PyUnicode_Check(value)
14164+
? json_encode_str(self, value)
14165+
: json_encode_dict_key_noinline(self, value);
14166+
}
14167+
else {
14168+
status = json_encode(self, value);
14169+
}
1415914170

1416014171
Py_DECREF(value);
1416114172
return status;

‎tests/unit/test_json.py‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2019,6 +2019,7 @@ def test_decode_dict_string_cache_ascii_only(self):
20192019
1,
20202020
1.5,
20212021
FruitInt.APPLE,
2022+
FruitStr.APPLE,
20222023
uuid.uuid4(),
20232024
datetime.datetime.now(),
20242025
datetime.date.today(),
@@ -2037,6 +2038,45 @@ def test_roundtrip_dict_key_types(self, key):
20372038
msg2 = msgspec.json.decode(sol, type=dict[type(key), int])
20382039
assert msg == msg2
20392040

2041+
def test_encode_dict_plain_enum_str_key(self):
2042+
# Regression: a plain `enum.Enum` with str values, used as a dict key.
2043+
# json.encode previously raised "Only dicts with str-like or
2044+
# number-like keys are supported", even though `dict[FruitStr, int]` is
2045+
# an accepted decode type and both msgpack and to_builtins handle it.
2046+
msg = {FruitStr.APPLE: 1, FruitStr.BANANA: 2}
2047+
2048+
res = msgspec.json.encode(msg)
2049+
assert res == b'{"apple":1,"banana":2}'
2050+
2051+
# JSON must agree with the msgpack encoder and with to_builtins.
2052+
expected = {"apple": 1, "banana": 2}
2053+
assert msgspec.msgpack.decode(msgspec.msgpack.encode(msg)) == expected
2054+
assert msgspec.to_builtins(msg, str_keys=True) == expected
2055+
2056+
# Full round-trip through the accepted decode type.
2057+
assert msgspec.json.decode(res, type=dict[FruitStr, int]) == msg
2058+
2059+
def test_encode_dict_plain_enum_str_key_nested(self):
2060+
# The same must hold when the dict is nested inside a Struct field.
2061+
class Rec(msgspec.Struct):
2062+
fruits: dict[FruitStr, int]
2063+
2064+
rec = Rec(fruits={FruitStr.APPLE: 5})
2065+
res = msgspec.json.encode(rec)
2066+
assert res == b'{"fruits":{"apple":5}}'
2067+
assert msgspec.json.decode(res, type=Rec) == rec
2068+
2069+
def test_encode_dict_plain_int_enum_key(self):
2070+
# A plain Enum with int values (not IntEnum) must keep round-tripping.
2071+
class IntVals(enum.Enum):
2072+
A = 1
2073+
B = 2
2074+
2075+
msg = {IntVals.A: 9, IntVals.B: 8}
2076+
res = msgspec.json.encode(msg)
2077+
assert res == b'{"1":9,"2":8}'
2078+
assert msgspec.json.decode(res, type=dict[IntVals, int]) == msg
2079+
20402080
def test_decode_dict_int_enum_key(self):
20412081
dec = msgspec.json.Decoder(dict[FruitInt, int])
20422082
assert dec.decode(b'{"-1": 10, "2": 20}') == {

0 commit comments

Comments
 (0)