python package for traversing data structures

45 Views Asked by At

I spent hours trying to figure out how to subclass and modify json.JSONEncoder to override how it encodes enum.Enum instances before learning that it literally wasn't possible, and that the only solution is to traverse my data structure and convert the enums myself in a pre-processing step.

Then I tried finding an existing solution to the problem of traversing arbitrary Python data structures, which I assumed was a very common and well-solved problem (either in Standard Library or PyPI), but I can't seem to find any.

1

There are 1 best solutions below

2
Grismar On

You didn't provide an example of what you're trying to achieve, but from the description, it appears you were after something like this:

from json import JSONEncoder, dumps
from enum import Enum


class Color(Enum):
    RED = 1
    GREEN = 2
    BLUE = 3


d = {'c': Color.RED, 'xs': [1, 3.14, Color.BLUE]}


class EnumNameJSONEncoder(JSONEncoder):
    def default(self, obj):
        if isinstance(obj, Enum):
            return obj.name
        return super().default(obj)


r = EnumNameJSONEncoder().encode(d)
print(dumps(r))

Output:

"{\"c\": \"RED\", \"xs\": [1, 3.14, \"BLUE\"]}"

If this doesn't solve your problem, you should provide a more specific example of what you're trying to achieve, what you tried and where you got stuck.

Edit: your comment suggested that the default() doesn't get called for enums - but that's not correct. For example, if you change the class like this:

class EnumNameJSONEncoder(JSONEncoder):
    def default(self, obj):
        print(f'getting called for {type(obj)}')
        if isinstance(obj, Enum):
            return obj.name
        return super().default(obj)

The output becomes:

getting called for <enum 'Color'>
getting called for <enum 'Color'>
"{\"c\": \"RED\", \"xs\": [1, 3.14, \"BLUE\"]}"