From 4511e4cdf70ded19048c7242c8e1d100155141f7 Mon Sep 17 00:00:00 2001
From: Louis Chauvet <louis.chauvet@free.fr>
Date: Sun, 22 Mar 2020 19:27:30 +0100
Subject: [PATCH] [base-storage] Custom type

---
 storage/jsonencoder.py | 34 ++++++++++++++++++++++++----------
 storage/objects.py     |  5 +++--
 2 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/storage/jsonencoder.py b/storage/jsonencoder.py
index a4dff78..7da7f7b 100644
--- a/storage/jsonencoder.py
+++ b/storage/jsonencoder.py
@@ -1,20 +1,34 @@
 import datetime
 import json
 
+data_type = "__data_type"
+content = "__content"
+
 
 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):
         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)):
-            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)
 
-
-def hook(dct):
-    if 'data_type' in dct:
-        if dct['data_type'] == "datetime.datetime":
-            return datetime.datetime.fromisoformat(dct['iso'])
-        elif dct['data_type'] == "datetime.timedelta":
-            return datetime.timedelta(seconds=dct['totalseconds'])
-    return dct
+    def hook(self, dct):
+        if data_type in dct:
+            for ty in self.custom.keys():
+                if str(ty) == dct[data_type]:
+                    return self.custom[ty][1](dct[content])
+            if dct[data_type] == "datetime.datetime":
+                return datetime.datetime.fromisoformat(dct['iso'])
+            elif dct[data_type] == "datetime.timedelta":
+                return datetime.timedelta(seconds=dct['totalseconds'])
+        return dct
diff --git a/storage/objects.py b/storage/objects.py
index 01856e1..ba9826e 100644
--- a/storage/objects.py
+++ b/storage/objects.py
@@ -8,17 +8,18 @@ class Objects:
     def __init__(self, path: str):
         self.path = os.path.abspath(path)
         os.makedirs(os.path.join(self.path, "objects"), exist_ok=True)
+        self.encoder = jsonencoder.Encoder
 
     def save_object(self, object_name, object_instance):
         """Save object into json 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):
         """Load object from json file"""
         if self.save_exists(object_name):
             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):
         """Check if json file exists"""