Search code examples
pythonpyyaml

Serializing custom classes with pyyaml


I would like to serialize a custom class (which I cannot modify or monkey patch) by specifying how such class instances should be serialized.

Here's the setup:

# some custom class which I cannot modify
class Custom:
    def __init__(self, a, b):
        self.a = a
        self.b = b

# data I want to serialize
data = [Custom(1, 2), Custom(101, 102)]

Here's how I would do it for JSON:

import json

# helper function to handle the custom class
def default(d):
    if isinstance(d, Custom):
        return dict(a=d.a, b=d.b)

print(json.dumps(data, default=default))
# expected and actual output: [{"a": 1, "b": 2}, {"a": 101, "b": 102}]

I'm struggling to find an equivalent solution for pyyaml:

import yaml

def yaml_equivalent_of_default():
    "YOUR SOLUTION GOES HERE"

print(yaml.dump(data))
# expected output:
# - a: 1
#   b: 2
# - a: 101
#   b: 102

I've tried different approaches mentioned in the pyyaml docs to no avail.


Solution

  • I believe this should work if the custom class members/fields can be pickled:

    import yaml
    
    class Custom:
        def __init__(self, a, b):
            self.a = a
            self.b = b
    
    data_to_serialize = [Custom(1, 2), Custom(101, 102)]
    
    def yaml_equivalent_of_default(dumper, data):
        dict_representation = data.__dict__
        node = dumper.represent_dict(dict_representation)
        return node
    
    yaml.add_representer(Custom, yaml_equivalent_of_default)
    
    print(yaml.dump(data_to_serialize))
    

    Output:

    - a: 1
      b: 2
    - a: 101
      b: 102
    

    The signature of the representer is def add_representer(data_type, representer, Dumper=Dumper), so it's possible to pass a Damper to it. All the available dumpers are ['BaseDumper', 'SafeDumper', 'Dumper'].