fix: json encoding

lib/src/prelude.py

@@ -39,31 +39,35 @@ def decode(s: bytes):
 
 
 class JSONEncoder(json.JSONEncoder):
-    def encode(self, o):
-        if isinstance(0, Json):
+    def default(self, o):
+        if isinstance(o, Json):
             return o.encode()
         elif isinstance(o, bytes):
-            return Json.encode(o.decode("base64"))
+            return b64encode(o).decode()
         elif isinstance(o, datetime):
-            return Json.encode(o.isoformat())
+            return o.isoformat()
         return json.JSONEncoder.encode(self, o)
 
 
 class JSONDecoder(json.JSONDecoder):
-    def decode(self, s):
-        o = json.JSONDecoder.decode(self, s)
-        if isinstance(o, bytes):
-            return o.encode("base64")
-
-        try:
-            return datetime.fromisoformat(o)
-        except:
-            pass
-
-        try:
-            return Json.decode(s)
-        except:
-            return o
+    def __init__(self, *args, **kwargs):
+        json.JSONDecoder.__init__(self, object_hook=self.object_hook, *args, **kwargs)
+
+    def object_hook(self, dct):
+        for k, v in dct.items():
+            if isinstance(v, str):
+                try:
+                    dct[k] = datetime.fromisoformat(v)
+                    continue
+                except:
+                    pass
+
+                try:
+                    dct[k] = b64decode(v)
+                    continue
+                except:
+                    pass
+        return dct
 
 
 class Json(Codec):
@@ -75,7 +79,8 @@ def encode(self) -> bytes:
 
     @classmethod
     def decode(cls, s: bytes):
-        return cls(**json.loads(s.decode(), cls=JSONDecoder))
+        x = json.loads(s.decode(), cls=JSONDecoder)
+        return cls(**x)
 
 
 class JsonObject(Json, dict):
@@ -88,7 +93,7 @@ def _store(x) -> int:
     elif isinstance(x, bytes):
         return ffi.memory.alloc(x).offset
     elif isinstance(x, dict) or isinstance(x, list):
-        return ffi.memory.alloc(json.dumps(x).encode()).offset
+        return ffi.memory.alloc(json.dumps(x, cls=JSONEncoder).encode()).offset
     elif isinstance(x, Codec):
         return ffi.memory.alloc(x.encode()).offset
     elif isinstance(x, ffi.memory.MemoryHandle):
@@ -114,7 +119,7 @@ def _load(t, x):
     elif t is bytes:
         return ffi.memory.bytes(mem)
     elif t is dict or t is list:
-        return json.loads(ffi.memory.string(mem))
+        return json.loads(ffi.memory.string(mem), cls=JSONDecoder)
     elif issubclass(t, Codec):
         return t.decode(ffi.memory.bytes(mem))
     elif t is ffi.memory.MemoryHandle:
@@ -171,15 +176,15 @@ def inner(*args):
 def input_json(t: Optional[type] = None):
     """Get input as JSON"""
     if t is not None:
-        return json.loads(input_str(), object_hook=lambda x: t(**x))
+        return json.loads(input_str(), object_hook=lambda x: t(**x), cls=JSONDecoder)
     return json.loads(input_str())
 
 
 def output_json(x):
     """Set JSON output"""
     if hasattr(x, "__dict__"):
         x = x.__dict__
-    output_str(json.dumps(x))
+    output_str(json.dumps(x, cls=JSONEncoder))
 
 
 def input(t: type = None):
@@ -192,7 +197,7 @@ def input(t: type = None):
     elif issubclass(t, Codec):
         return t.decode(input_bytes())
     elif t is dict or t is list:
-        return json.loads(input_str())
+        return json.loads(input_str(), cls=JSONDecoder)
     elif issubclass(t, Enum):
         return t(input_str())
     else: