tinydb-yaml-mwe.py

#!/usr/bin/env python3
from tinydb import TinyDB
from tinydb.storages import Storage, touch

import yaml


class TestYamlStorage1(Storage):
    """
    Store the data in a YAML file.
    Written following the JSONStorage class in TinyDB as closely as possible.
    """

    def __init__(self, path, **kwargs):
        """
        Create a new instance.
        Also creates the storage file, if it doesn't exist.
        :param path: Where to store the YAML data.
        :type path: str
        """

        super().__init__()
        touch(path)  # Create file if not exists
        self.kwargs = kwargs
        self._handle = open(path, 'r+')

    def close(self):
        self._handle.close()

    def read(self):
        # Get the file size
        self._handle.seek(0, 2)
        size = self._handle.tell()

        if not size:
            # File is empty
            return None
        else:
            self._handle.seek(0)
            loaded = yaml.load(self._handle)
            print('loaded: {}'.format(loaded))
            return loaded

    def write(self, data):
        # print(type(data))
        # print('data: {}'.format(data))
        self._handle.seek(0)
        # serialized = json.dumps(data, **self.kwargs)
        # print(self.kwargs)
        serialized = yaml.dump(data, **self.kwargs)
        # print('serialized: {}'.format(serialized))
        self._handle.write(serialized)
        self._handle.flush()
        self._handle.truncate()


class TestYamlStorage2(Storage):
    """
    Store the data in a YAML file.
    Written following the example at http://tinydb.readthedocs.io/en/latest/extend.html#write-a-custom-storage
    """
    def __init__(self, filename):  # (1)
        super().__init__()
        self.filename = filename
        touch(filename)

    def read(self):
        with open(self.filename) as handle:
            try:
                # data = ruamel.yaml.load(handle, ruamel.yaml.RoundTripLoader)
                # data = yaml.safe_load(handle.read())
                data = yaml.load(handle.read())
                return data
            except yaml.YAMLError:
                return None  # (3)

    def write(self, data):
        print('writing data: {}'.format(data))
        with open(self.filename, 'w') as handle:
                yaml.dump(data, handle)
                # ruamel.yaml.dump(data, stream=handle, Dumper=ruamel.yaml.RoundTripDumper)

    def close(self):  # (4)
        pass

def main():

    import os

    # flush existing file for consistency
    os.remove('db.yaml')

    db = TinyDB('db.yaml', storage=TestYamlStorage2)
    dicts = [
        dict(name='Homer', age=38),
        dict(name='Marge', age=34),
        dict(name='Bart', age=10)
    ]

    # this works as expected
    # db.insert_multiple(dicts)

    # this doesn't
    db.insert(dicts[0])
    db.insert(dicts[1])
    print(db.all())

if __name__ == '__main__':
    main()