[tor-commits] [sbws/master] chg: json: Create custom JSON encoder/decoder

[juga at torproject.org]

commit f19738f1c24316399c0928f62130ec133cee5db0
Author: juga0 <juga at riseup.net>
Date:   Sat Mar 14 16:52:21 2020 +0000

    chg: json: Create custom JSON encoder/decoder
    
    to be able to serialize/deserailize datetime in the state file.
---
 sbws/util/json.py            | 44 ++++++++++++++++++++++++++++++++++++
 tests/unit/util/test_json.py | 54 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 98 insertions(+)

diff --git a/sbws/util/json.py b/sbws/util/json.py
new file mode 100644
index 0000000..f4f8130
--- /dev/null
+++ b/sbws/util/json.py
@@ -0,0 +1,44 @@
+"""JSON custom serializers and deserializers."""
+import datetime
+import json
+
+from .timestamps import DateTimeSeq, DateTimeIntSeq
+
+
+class CustomEncoder(json.JSONEncoder):
+    """JSONEncoder that serializes datetime to ISO 8601 string."""
+
+    def default(self, obj):
+        if isinstance(obj, DateTimeSeq) or isinstance(obj, DateTimeIntSeq):
+            return [self.default(i) for i in obj.list()]
+        if isinstance(obj, datetime.datetime):
+            return obj.replace(microsecond=0).isoformat()
+        else:
+            return super().default(obj)
+
+
+class CustomDecoder(json.JSONDecoder):
+    """JSONDecoder that deserializes ISO 8601 string to datetime."""
+
+    def decode(self, s, **kwargs):
+        decoded = super().decode(s, **kwargs)
+        return self.process(decoded)
+
+    def process(self, obj):
+        if isinstance(obj, list) and obj:
+            return [self.process(item) for item in obj]
+        if isinstance(obj, dict):
+            return {key: self.process(value) for key, value in obj.items()}
+        if isinstance(obj, str):
+            try:
+                return datetime.datetime.strptime(obj, "%Y-%m-%dT%H:%M:%S")
+            except ValueError:
+                try:
+                    datetime.datetime.strptime(
+                        obj, "%Y-%m-%dT%H:%M:%S.%f"
+                    ).replace(microsecond=0)
+                except ValueError:
+                    pass
+            except TypeError:
+                pass
+        return obj
diff --git a/tests/unit/util/test_json.py b/tests/unit/util/test_json.py
new file mode 100644
index 0000000..05e0bf3
--- /dev/null
+++ b/tests/unit/util/test_json.py
@@ -0,0 +1,54 @@
+"""json.py unit tests."""
+import json
+
+from sbws.util.json import CustomDecoder, CustomEncoder
+
+STATE = """{
+    "min_perc_reached": null,
+    "recent_consensus_count": [
+        "2020-03-04T10:00:00",
+        "2020-03-05T10:00:00",
+        "2020-03-06T10:00:00"
+    ],
+    "recent_measurement_attempt": [
+        [
+            "2020-03-04T10:00:00",
+            2
+        ],
+        [
+            "2020-03-05T10:00:00",
+            2
+        ],
+        [
+            "2020-03-06T10:00:00",
+            2
+        ]
+    ],
+    "recent_priority_list": [
+        "2020-03-04T10:00:00",
+        "2020-03-05T10:00:00",
+        "2020-03-06T10:00:00"
+    ],
+    "recent_priority_relay": [
+        [
+            "2020-03-04T10:00:00",
+            2
+        ],
+        [
+            "2020-03-05T10:00:00",
+            2
+        ],
+        [
+            "2020-03-06T10:00:00",
+            2
+        ]
+    ],
+    "scanner_started": "2020-03-14T16:15:22",
+    "uuid": "x"
+}"""
+
+
+def test_decode_encode_roundtrip():
+    d = json.loads(STATE, cls=CustomDecoder)
+    s = json.dumps(d, cls=CustomEncoder, indent=4, sort_keys=True)
+    assert s == STATE