add support for datetime and regex in json_util PYTHON-76

Mike Dirolf · Mike Dirolf · commit e4d36aea11e0 · 2009-12-09T11:08:16.000-05:00
diff --git a/pymongo/json_util.py b/pymongo/json_util.py
@@ -30,28 +30,63 @@
 
 >>> json.loads(..., object_hook=json_util.object_hook)
 
-Currently this only handles special encoding and decoding for
-:class:`~pymongo.objectid.ObjectId` and :class:`~pymongo.dbref.DBRef`
+Currently this does not handle special encoding and decoding for
+:class:`~pymongo.binary.Binary` and :class:`~pymongo.code.Code`
 instances.
+
+.. versionchanged:: 1.1.2+
+   Added support for encoding/decoding datetimes and regular expressions.
 """
 
+import datetime
+import calendar
+import re
+
 from objectid import ObjectId
 from dbref import DBRef
 
-# TODO support other types, like Binary, Code, datetime & regex
+# TODO support Binary and Code
 # Binary and Code are tricky because they subclass str so json thinks it can
 # handle them. Not sure what the proper way to get around this is...
+#
+# One option is to just add some other method that users need to call _before_
+# calling json.dumps or json.loads. That is pretty terrible though...
+
+# TODO share this with bson.py?
+_RE_TYPE = type(re.compile("foo"))
 
 def object_hook(dct):
     if "$oid" in dct:
         return ObjectId(str(dct["$oid"]))
     if "$ref" in dct:
         return DBRef(dct["$ref"], dct["$id"], dct.get("$db", None))
+    if "$date" in dct:
+        return datetime.datetime.utcfromtimestamp(float(dct["$date"]) / 1000.0)
+    if "$regex" in dct:
+        flags = 0
+        if "i" in dct["$options"]:
+            flags |= re.IGNORECASE
+        if "m" in dct["$options"]:
+            flags |= re.MULTILINE
+        return re.compile(dct["$regex"], flags)
     return dct
 
 def default(obj):
     if isinstance(obj, ObjectId):
         return {"$oid": str(obj)}
     if isinstance(obj, DBRef):
         return obj.as_doc()
+    if isinstance(obj, datetime.datetime):
+        # TODO share this code w/ bson.py?
+        millis = int(calendar.timegm(obj.timetuple()) * 1000 +
+                     obj.microsecond / 1000)
+        return {"$date": str(millis)}
+    if isinstance(obj, _RE_TYPE):
+        flags = ""
+        if obj.flags & re.IGNORECASE:
+            flags += "i"
+        if obj.flags & re.MULTILINE:
+            flags += "m"
+        return {"$regex": obj.pattern,
+                "$options": flags}
     raise TypeError("%r is not JSON serializable" % obj)
diff --git a/test/test_json_util.py b/test/test_json_util.py
@@ -15,6 +15,8 @@
 """Test some utilities for working with JSON and PyMongo."""
 
 import unittest
+import datetime
+import re
 import sys
 json_lib = True
 try:
@@ -39,12 +41,12 @@ def setUp(self):
         if not json_lib:
             raise SkipTest()
 
-    def round_trip(self, doc):
-        def round_tripped(doc):
-            return json.loads(json.dumps(doc, default=default),
-                              object_hook=object_hook)
+    def round_tripped(self, doc):
+        return json.loads(json.dumps(doc, default=default),
+                          object_hook=object_hook)
 
-        self.assertEqual(doc, round_tripped(doc))
+    def round_trip(self, doc):
+        self.assertEqual(doc, self.round_tripped(doc))
 
     def test_basic(self):
         self.round_trip({"hello": "world"})
@@ -56,6 +58,16 @@ def test_dbref(self):
         self.round_trip({"ref": DBRef("foo", 5)})
         self.round_trip({"ref": DBRef("foo", 5, "db")})
 
+    def test_datetime(self):
+        # only millis, not micros
+        self.round_trip({"date": datetime.datetime(2009, 12, 9, 15,
+                                                   49, 45, 191000)})
+
+    def test_regex(self):
+        res = self.round_tripped({"r": re.compile("a*b", re.IGNORECASE)})["r"]
+        self.assertEqual("a*b", res.pattern)
+        self.assertEqual(re.IGNORECASE, res.flags)
+
 
 if __name__ == "__main__":
     unittest.main()
