From 9c354474cb6becf683e56c29eff49c9fc040f84e Mon Sep 17 00:00:00 2001 From: Lucas Cardozo Date: Tue, 23 Apr 2019 17:22:26 -0300 Subject: [PATCH] Adding support for yaml serialization --- pickledb.py | 66 ++++++++++++++++++++++++++++++----------------------- tests.py | 10 ++++---- 2 files changed, 44 insertions(+), 32 deletions(-) diff --git a/pickledb.py b/pickledb.py index 40fbece..6d3c6bb 100644 --- a/pickledb.py +++ b/pickledb.py @@ -4,53 +4,64 @@ # Copyright 2018 Harrison Erd # -# Redistribution and use in source and binary forms, with or without +# Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions are met: # -# 1. Redistributions of source code must retain the above copyright notice, +# 1. Redistributions of source code must retain the above copyright notice, # this list of conditions and the following disclaimer. # -# 2. Redistributions in binary form must reproduce the above copyright notice, -# this list of conditions and the following disclaimer in the documentation +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation # and/or other materials provided with the distribution. # -# 3. Neither the name of the copyright holder nor the names of its -# contributors may be used to endorse or promote products derived from this +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from this # software without specific prior written permission. # -# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS -# IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, -# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR -# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR -# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, -# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, -# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; -# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, -# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR -# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS +# IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, +# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; +# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR +# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, # EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. import sys import os import signal -import json +from functools import partial from threading import Thread -def load(location, auto_dump, sig=True): +def load(location, auto_dump, sig=True, serializer="auto"): '''Return a pickledb object. location is the path to the json file.''' - return PickleDB(location, auto_dump, sig) + return PickleDB(location, auto_dump, sig, serializer) class PickleDB(object): key_string_error = TypeError('Key/name must be a string!') - def __init__(self, location, auto_dump, sig): + def __init__(self, location, auto_dump, sig, serializer): '''Creates a database object and loads the data from the location path. If the file does not exist it will be created on the first update. ''' + if serializer == 'auto': + serializer = location.split('.')[-1] + + if serializer == 'json': + import json + self.serializer = json + elif serializer == 'yaml' or serializer == 'yml': + import yaml + self.serializer = yaml + self.serializer.load = partial(yaml.load, Loader=yaml.SafeLoader) + self.load(location, auto_dump) self.dthread = None if sig: @@ -89,20 +100,20 @@ def load(self, location, auto_dump): def dump(self): '''Force dump memory db to file''' - json.dump(self.db, open(self.loco, 'wt')) + self.serializer.dump(self.db, open(self.loco, 'wt')) self.dthread = Thread( - target=json.dump, + target=self.serializer.dump, args=(self.db, open(self.loco, 'wt'))) self.dthread.start() self.dthread.join() return True - + def _loaddb(self): - '''Load or reload the json info from the file''' - self.db = json.load(open(self.loco, 'rt')) + '''Load or reload the info from the file''' + self.db = self.serializer.load(open(self.loco, 'rt')) def _autodumpdb(self): - '''Write/save the json dump into the file if auto_dump is enabled''' + '''Write/save the dump into the file if auto_dump is enabled''' if self.auto_dump: self.dump() @@ -132,7 +143,7 @@ def exists(self, key): def rem(self, key): '''Delete a key''' - if not key in self.db: # return False instead of an exception + if key not in self.db: # return False instead of an exception return False del self.db[key] self._autodumpdb() @@ -279,4 +290,3 @@ def deldb(self): self.db = {} self._autodumpdb() return True - diff --git a/tests.py b/tests.py index 7e5141b..e8238a3 100644 --- a/tests.py +++ b/tests.py @@ -9,8 +9,10 @@ class TestClass(object): db = pickledb.load('tests.db', auto_dump=False) def test_load(self): - x = pickledb.load('x.db', auto_dump=False) - assert x is not None + jsondb = pickledb.load('x.json', auto_dump=False, serializer='json') + yamldb = pickledb.load('x.yaml', auto_dump=False, serializer='yaml') + assert jsondb is not None + assert yamldb is not None def test_sugar_get(self): self.db.db["foo"] = "bar" @@ -75,14 +77,14 @@ def test_lexists(self): self.db.ladd('list', 'value') x = self.db.lexists('list', 'value') assert x is True - self.db.lrem('list') + self.db.lremlist('list') def test_not_lexists(self): self.db.lcreate('list') self.db.ladd('list', 'value') x = self.db.lexists('list', 'not_value') assert x is False - self.db.lrem('list') + self.db.lremlist('list') def test_dexists(self): self.db.dcreate('dict')