Skip to content

Instantly share code, notes, and snippets.

@fndari
Created May 17, 2016 15:29
Show Gist options
  • Save fndari/f6885b336359370d4afbd5a9340316cb to your computer and use it in GitHub Desktop.
Save fndari/f6885b336359370d4afbd5a9340316cb to your computer and use it in GitHub Desktop.
#!/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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment