Python Console cannot find associated class when loading from yaml file

Question

I would like to save an instance of a class in a YAML File (to make it more human readable).

Consider this example code:

import tempfile
from pathlib import Path

import yaml


class MyClass:
    def __init__(self, attribute):
        self.attribute = attribute

if __name__=='__main__':

    my_instance = MyClass(attribute="Hello World")

    with tempfile.TemporaryDirectory() as tmpdir:
        # save file
        with open(Path(tmpdir) / 'my_save.yaml','w') as file:
            yaml.dump(my_instance,file)

        # load file
        with open(Path(tmpdir) / 'my_save.yaml','r') as file:
            my_instance_loaded = yaml.load(file,Loader= yaml.UnsafeLoader)

        print(my_instance_loaded.attribute)

This code does work without hiccups if I run the script via the commandline.

But if run the same code in the Python console, I get the following error:

yaml.constructor.ConstructorError: while constructing a Python object cannot find 'MyClass' in the module '__main__'

While this is not an acute problem, my approach seems to be error prone. Especially if I load the yaml file in other parts of my program.

Is there a way to tell the yaml loader just to go with the name of the class (e.g.MyClass) instead looking up the whole "object-Path" (e.g.__main__.MyClass)?

Answer 1

If you don't mind modifying your objects so that they're registered for PyYAML, you can make MyClass inherit from yaml.YAMLObject, like they instruct in the documentation

i.e:

import yaml

class MyClass(yaml.YAMLObject):
    yaml_tag = '!MyClass'

    def __init__(self, attribute):
        self.attribute = attribute

my_instance = MyClass(attribute="Hello World")

dumped = yaml.dump(my_instance)

print(dumped)
> !MyClass
> attribute: Hello World

loaded = yaml.load(dumped)

print(loaded.__dict__)
> {'attribute': 'Hello World'}

print(my_instance.__dict__)
> {'attribute': 'Hello World'}

Alternatively, you could manually register constructors and representers for your objects like this:

import yaml

class MyClass:
    def __init__(self, attribute):
        self.attribute = attribute

# This maps objects attributes to a dictionary and a `yaml_tag` to it
def myclass_representer(dumper, data):
    return dumper.represent_mapping("!MyClass", {"attribute": data.attribute})

# This maps a `MappingNode` to the construction of a class. Note that we construct a mapping because before we represented a mapping
def myclass_constructor(loader, node):
    value = loader.construct_mapping(node)
    return MyClass(**value)

yaml.add_representer(MyClass, myclass_representer)
yaml.add_constructor('!MyClass', myclass_constructor)

my_instance = MyClass(attribute="Hello World")

dumped = yaml.dump(my_instance)

print(dumped)
> !MyClass
> attribute: Hello World

loaded = yaml.load(dumped)

print(loaded.__dict__)
> {'attribute': 'Hello World'}

print(my_instance.__dict__)
> {'attribute': 'Hello World'}