PDBA/bot-base
8
2
Fork 0
Code Issues Pull Requests Releases Wiki Activity

[base-storage] Custom type

Browse Source
This commit is contained in:
Louis Chauvet 2020-03-22 19:27:30 +01:00
parent 45848dc2f8
commit 4511e4cdf7
Signed by: fomys
GPG Key ID: 1ECA046A9615ABA0
2 changed files with 27 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
storage/jsonencoder.py
View File

@ -1,20 +1,34 @@
import datetime import datetime
import json import json
data_type = "__data_type"
content = "__content"
class Encoder(json.JSONEncoder): class Encoder(json.JSONEncoder):
def __init__(self):
super().__init__()
self.custom = {}
def register(self, type_, encode, decode):
self.custom.update({type_: (encode, decode)})
def default(self, obj): def default(self, obj):
if isinstance(obj, (datetime.datetime)): if isinstance(obj, (datetime.datetime)):
return {'data_type':'datetime.datetime', 'iso':obj.isoformat()} return {data_type: 'datetime.datetime', 'iso': obj.isoformat()}
if isinstance(obj, (datetime.timedelta)): if isinstance(obj, (datetime.timedelta)):
return {'data_type':'datetime.timedelta', 'totalseconds':obj.total_seconds()} return {data_type: 'datetime.timedelta', 'totalseconds': obj.total_seconds()}
if isinstance(obj, tuple(self.custom.keys())):
return {data_type: f'{type(obj)}', content: self.custom[type(obj)][0](obj)}
return json.JSONEncoder.default(self, obj) return json.JSONEncoder.default(self, obj)
def hook(self, dct):
def hook(dct): if data_type in dct:
if 'data_type' in dct: for ty in self.custom.keys():
if dct['data_type'] == "datetime.datetime": if str(ty) == dct[data_type]:
return datetime.datetime.fromisoformat(dct['iso']) return self.custom[ty][1](dct[content])
elif dct['data_type'] == "datetime.timedelta": if dct[data_type] == "datetime.datetime":
return datetime.timedelta(seconds=dct['totalseconds']) return datetime.datetime.fromisoformat(dct['iso'])
return dct elif dct[data_type] == "datetime.timedelta":
return datetime.timedelta(seconds=dct['totalseconds'])
return dct

5
storage/objects.py
View File

@ -8,17 +8,18 @@ class Objects:
def __init__(self, path: str): def __init__(self, path: str):
self.path = os.path.abspath(path) self.path = os.path.abspath(path)
os.makedirs(os.path.join(self.path, "objects"), exist_ok=True) os.makedirs(os.path.join(self.path, "objects"), exist_ok=True)
self.encoder = jsonencoder.Encoder
def save_object(self, object_name, object_instance): def save_object(self, object_name, object_instance):
"""Save object into json file""" """Save object into json file"""
with open(os.path.join(self.path, "objects", object_name + ".json"), "w") as file: with open(os.path.join(self.path, "objects", object_name + ".json"), "w") as file:
json.dump([object_instance], file, cls=jsonencoder.Encoder) json.dump(object_instance, file, cls=self.encoder)
def load_object(self, object_name): def load_object(self, object_name):
"""Load object from json file""" """Load object from json file"""
if self.save_exists(object_name): if self.save_exists(object_name):
with open(os.path.join(self.path, "objects", object_name + ".json"), "r") as f: with open(os.path.join(self.path, "objects", object_name + ".json"), "r") as f:
return json.load(f, object_hook=jsonencoder.hook)[0] return json.load(f, object_hook=self.encoder.hook)
def save_exists(self, object_name): def save_exists(self, object_name):
"""Check if json file exists""" """Check if json file exists"""