From 3e6d3f3cbe23a0a9be3c86d4b53cf8e1ea5e5f5b Mon Sep 17 00:00:00 2001
From: Mike Dirolf <mike@10gen.com>
Date: Wed, 7 Jan 2009 17:10:00 -0500
Subject: [PATCH] support for dates

---
 bson.py        | 16 +++++++++++++++-
 test/qcheck.py | 12 ++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/bson.py b/bson.py
index 4cbf306a5..f190d2a01 100644
--- a/bson.py
+++ b/bson.py
@@ -7,6 +7,8 @@ import types
 import struct
 import random
 import re
+import datetime
+import time
 
 from test import test_data, qcheck
 
@@ -186,6 +188,10 @@ def _get_binary(data):
 def _get_boolean(data):
     return (data[0] == "\x01", data[1:])
 
+def _get_date(data):
+    seconds = float(struct.unpack("<q", data[:8])[0]) / 1000.0
+    return (datetime.datetime.fromtimestamp(seconds), data[8:])
+
 def _get_null(data):
     return (None, data)
 
@@ -198,7 +204,7 @@ _element_getter = {
     "\x06": _get_null, # undefined
 #    "\x07": _get_oid,
     "\x08": _get_boolean,
-#    "\x09": _get_date,
+    "\x09": _get_date,
     "\x0A": _get_null,
 #    "\x0B": _get_regex,
 #    "\x0C": _get_ref,
@@ -229,6 +235,9 @@ def _document_to_dict(data):
 def _int_to_bson(int):
     return struct.pack("<i", int)
 
+def _int_64_to_bson(int):
+    return struct.pack("<q", int)
+
 def _value_to_bson(value):
     if isinstance(value, types.FloatType):
         return ("\x01", struct.pack("<d", value))
@@ -247,6 +256,9 @@ def _value_to_bson(value):
         if value:
             return ("\x08", "\x01")
         return ("\x08", "\x00")
+    if isinstance(value, datetime.datetime):
+        millis = int(time.mktime(value.timetuple()) * 1000 + value.microsecond / 1000)
+        return ("\x09", _int_64_to_bson(millis))
     if isinstance(value, types.NoneType):
         return ("\x0A", "")
     if isinstance(value, types.IntType):
@@ -374,6 +386,8 @@ class TestBSON(unittest.TestCase):
                          "\x13\x00\x00\x00\x05\x74\x65\x73\x74\x00\x04\x00\x00\x00\x74\x65\x73\x74\x00")
         self.assertEqual(BSON.from_dict({"test": None}),
                          "\x0B\x00\x00\x00\x0A\x74\x65\x73\x74\x00\x00")
+        self.assertEqual(BSON.from_dict({"date": datetime.datetime(2007, 1, 7, 19, 30, 11)}),
+                         "\x13\x00\x00\x00\x09\x64\x61\x74\x65\x00\x38\xBE\x1C\xFF\x0F\x01\x00\x00\x00")
 
     def test_from_then_to_dict(self):
         def helper(dict):
diff --git a/test/qcheck.py b/test/qcheck.py
index a01840a17..54f72cba7 100644
--- a/test/qcheck.py
+++ b/test/qcheck.py
@@ -1,6 +1,7 @@
 import random
 import sys
 import traceback
+import datetime
 
 gen_target = 100
 examples = 5
@@ -44,6 +45,16 @@ def gen_unicode(gen_length):
 def gen_list(generator, gen_length):
     return lambda: [generator() for _ in range(gen_length())]
 
+# This is a little Mongo specific...
+def gen_datetime():
+    return lambda: datetime.datetime(random.randint(1970, 2037),
+                                     random.randint(1, 12),
+                                     random.randint(1, 28),
+                                     random.randint(0, 23),
+                                     random.randint(0, 59),
+                                     random.randint(0, 59),
+                                     random.randint(0, 999) * 1000)
+
 def gen_dict(gen_key, gen_value, gen_length):
     def a_dict(gen_key, gen_value, length):
         result = {}
@@ -58,6 +69,7 @@ def gen_mongo_value(depth):
                gen_int(),
                gen_float(),
                gen_boolean(),
+               gen_datetime(),
                lift(None),]
     if depth > 0:
         choices.append(gen_mongo_list(depth))