Encoding Python Enum to JSON

问题

I have a dictionary where some of the keys are Enum instances (subclasses of enum.Enum). I am attempting to encode the dictionary into a JSON string using a custom JSON Encoder class as per the documentation. All I want is to have the keys in the outputted JSON be the strings of the Enum names. For example { TestEnum.one : somevalue } would be encoded to { "one" : somevalue }.

I have written a simple test case, shown below, which I have tested in a clean virtualenv:

import json

from enum import Enum

class TestEnum(Enum):
    one = "first"
    two = "second"
    three = "third"

class TestEncoder(json.JSONEncoder):
    """ Custom encoder class """

    def default(self, obj):

        print("Default method called!")

        if isinstance(obj, TestEnum):
            print("Seen TestEnum!")
            return obj.name

        return json.JSONEncoder.default(self, obj)

def encode_enum(obj):
    """ Custom encoder method """

    if isinstance(obj, TestEnum):
        return obj.name
    else:
        raise TypeError("Don't know how to decode this")

if __name__ == "__main__":

    test = {TestEnum.one : "This",
            TestEnum.two : "should",
            TestEnum.three : "work!"}

    # Test dumps with Encoder method
    #print("Test with encoder method:")
    #result = json.dumps(test, default=encode_enum)
    #print(result)

    # Test dumps with Encoder Class
    print("Test with encoder class:")
    result = json.dumps(test, cls=TestEncoder)
    print(result)

I cannot successfully encode the dictionary (using Python 3.6.1). I continually get TypeError: keys must be a string errors and the default method of my custom encoder instance (supplied via the cls argument of the json.dumps method) never seems to be called? I have also attempted to supply a custom encoding method via the default argument of the json.dumps method, but again this is never triggered.

I have seen solutions involving the IntEnum class, but I need the values of the Enum to be strings. I have also seen this answer which discusses an issue related to an Enum which inherits from another class. However, my enums inherit from the base enum.Enum class only and correctly respond to isinstance calls?

Both the custom class and the method produce a TypeError when supplied to the json.dumps method. Typical output is shown below:

$ python3 enum_test.py

Test with encoder class
Traceback (most recent call last):
  File "enum_test.py", line 59, in <module>
    result = json.dumps(test, cls=TestEncoder)
  File "/usr/lib64/python3.6/json/__init__.py", line 238, in dumps
    **kw).encode(obj)
  File "/usr/lib64/python3.6/json/encoder.py", line 199, in encode
    chunks = self.iterencode(o, _one_shot=True)
  File "/usr/lib64/python3.6/json/encoder.py", line 257, in iterencode
    return _iterencode(o, 0)
TypeError: keys must be a string

I presume the issue is that the encode method of the JSONEncoder class assumes that it knows how to serialise the Enum class (because one of the if statements in the iterencode method is triggered) and so never calls the custom default methods and ends failing to serialise the Enum?

Any help would be greatly appreciated!

回答1:

You can't use anything but strings as keys in dictionaries you want to convert to JSON. The encoder doesn't give you any other options; the default hook is only called for values of unknown type, never for keys.

Convert your keys to strings up front:

def convert_keys(obj, convert=str):
    if isinstance(obj, list):
        return [convert_keys(i, convert) for i in obj]
    if not isinstance(obj, dict):
        return obj
    return {convert(k): convert_keys(v, convert) for k, v in obj.items()}

json.dumps(convert_keys(test))

This recursively handles your dictionary keys. Note that I included a hook; you can then choose how to convert enumeration values to strings:

def enum_names(key):
    if isinstance(key, TestEnum):
        return key.name
    return str(key)

json.dumps(convert_keys(test, enum_names))

You can use the same function to reverse the process when loading from JSON:

def names_to_enum(key):
    try:
        return TestEnum[key]
    except KeyError:
        return key

convert_keys(json.loads(json_data), names_to_enum)

Demo:

>>> def enum_names(key):
...     if isinstance(key, TestEnum):
...         return key.name
...     return str(key)
...
>>> json_data = json.dumps(convert_keys(test, enum_names))
>>> json_data
'{"one": "This", "two": "should", "three": "work!"}'
>>> def names_to_enum(key):
...     try:
...         return TestEnum[key]
...     except KeyError:
...         return key
...
>>> convert_keys(json.loads(json_data), names_to_enum)
{<TestEnum.one: 'first'>: 'This', <TestEnum.two: 'second'>: 'should', <TestEnum.three: 'third'>: 'work!'}

来源：https://stackoverflow.com/questions/43854335/encoding-python-enum-to-json

标签

python

json

enums

python-3.6