Writing a Custom Storage for YAML support

[TheNewThinkTank]

Following official guidelines, https://tinydb.readthedocs.io/en/latest/extend.html , I attempt to add YAML support to my TinyDB project as follows (my complete script):

import yaml  # type: ignore

from tinydb import TinyDB  # type: ignore


class YAMLStorage(Storage):
    def __init__(self, filename):
        self.filename = filename

    def read(self):
        with open(self.filename) as handle:
            try:
                data = yaml.safe_load(handle.read())
                return data
            except yaml.YAMLError:
                return None

    def write(self, data):
        with open(self.filename, 'w+') as handle:
            yaml.dump(data, handle)

    def close(self):
        pass


if __name__ == "__main__":
    db = TinyDB('db.yml', storage=YAMLStorage)

And receive following name error (unknown base class):

    class YAMLStorage(Storage):
                      ^^^^^^^
NameError: name 'Storage' is not defined

Is there some additional import I am missing perhaps?

[TheNewThinkTank]

For reference, here is the specific commit, with GitHub Actions workflow failing with the above mentioned error: TheNewThinkTank/fitness-tracker@8bf84d3

[TheNewThinkTank]

Solved by importing Storage from tinydb.storages

[msiemens]

Glad to hear to were able to solve this 🙂

[TheNewThinkTank]

Thanks, @msiemens , works like a charm now, and I swapped my project to using YAML-based TinyDB only