Merge pull request #729 from esoteric-ephemera/master

Add support for extended JSON to MontyDecoder

src/monty/json.py

@@ -29,8 +29,10 @@
 
 try:
     import bson
+    from bson import json_util
 except ImportError:
     bson = None
+    json_util = None
 
 try:
     import orjson
@@ -862,7 +864,11 @@ def decode(self, s):
         :param s: string
         :return: Object.
         """
-        if orjson is not None:
+        if bson is not None:
+            # need to pass `json_options` to ensure that datetimes are not
+            # converted by BSON
+            d = json_util.loads(s, json_options=json_util.JSONOptions(tz_aware=True))
+        elif orjson is not None:
             try:
                 d = orjson.loads(s)
             except orjson.JSONDecodeError:

tests/test_json.py

@@ -1068,3 +1068,31 @@ def test_enum(self):
         assert d_ == {"v": "value_a"}
         na2 = EnumAsDict.from_dict(d_)
         assert na2 == na1
+
+    @pytest.mark.skipif(ObjectId is None, reason="bson not present")
+    def test_extended_json(self):
+        from bson import json_util
+
+        ext_json_dict = {
+            "datetime": datetime.datetime.now(datetime.timezone.utc),
+            "NaN": float("NaN"),
+            "infinity": float("inf"),
+            "-infinity": -float("inf"),
+        }
+        ext_json_str = json_util.dumps(ext_json_dict)
+
+        not_serialized = json.loads(ext_json_str)
+        assert all(isinstance(v, dict) for v in not_serialized.values())
+
+        reserialized = MontyDecoder().decode(ext_json_str)
+        for k, v in ext_json_dict.items():
+            if k == "datetime":
+                # BSON's json_util only saves datetimes up to microseconds
+                assert reserialized[k].timestamp() == pytest.approx(
+                    v.timestamp(), abs=1e-3
+                )
+            elif k == "NaN":
+                assert np.isnan(reserialized[k])
+            else:
+                assert v == reserialized[k]
+            assert not isinstance(reserialized[k], dict)