PYTHON-707 Add a BSONInt64 type.

Luke Lovett · Luke Lovett · commit 5136bb723cc5 · 2014-07-18T14:59:21.000Z
diff --git a/bson/__init__.py b/bson/__init__.py
@@ -24,6 +24,7 @@
 
 from bson.binary import (Binary, OLD_UUID_SUBTYPE,
                          JAVA_LEGACY, CSHARP_LEGACY)
+from bson.bsonint64 import BSONInt64
 from bson.code import Code
 from bson.dbref import DBRef
 from bson.errors import (InvalidBSON,
@@ -280,7 +281,7 @@ def _get_timestamp(
 def _get_long(data, position, as_class, tz_aware, uuid_subtype, compile_re):
     # Have to cast to long (for python 2.x); on 32-bit unpack may return
     # an int.
-    value = long(struct.unpack("<q", data[position:position + 8])[0])
+    value = BSONInt64(struct.unpack("<q", data[position:position + 8])[0])
     position += 8
     return value, position
 
@@ -416,6 +417,8 @@ def _element_to_bson(key, value, check_keys, uuid_subtype):
         return BSONBOO + name + ONE
     if value is False:
         return BSONBOO + name + ZERO
+    if isinstance(value, BSONInt64):
+        return BSONLON + name + struct.pack("<q", value)
     if isinstance(value, int):
         # TODO this is an ugly way to check for this...
         if value > MAX_INT64 or value < MIN_INT64:
diff --git a/bson/_cbsonmodule.c b/bson/_cbsonmodule.c
@@ -50,6 +50,7 @@ struct module_state {
     PyObject* MaxKey;
     PyObject* UTC;
     PyTypeObject* REType;
+    PyObject* BSONInt64;
 };
 
 /* The Py_TYPE macro was introduced in CPython 2.6 */
@@ -330,6 +331,7 @@ static int _load_python_objects(PyObject* module) {
         _load_object(&state->UTC, "bson.tz_util", "utc") ||
         _load_object(&state->RECompile, "re", "compile") ||
         _load_object(&state->Regex, "bson.regex", "Regex") ||
+        _load_object(&state->BSONInt64, "bson.bsonint64", "BSONInt64") ||
         _load_object(&state->UUID, "uuid", "UUID")) {
         return 1;
     }
@@ -2028,16 +2030,17 @@ static PyObject* get_value(PyObject* self, const char* buffer, unsigned* positio
         }
     case 18:
         {
+            PyObject* bsonint64_type = _get_object(state->BSONInt64, "bson.bsonint64", "BSONInt64");
+            if (!bsonint64_type)
+                goto invalid;
             long long ll;
             if (max < 8) {
                 goto invalid;
             }
             memcpy(&ll, buffer + *position, 8);
-            value = PyLong_FromLongLong(ll);
-            if (!value) {
-                goto invalid;
-            }
+            value = PyObject_CallFunction(bsonint64_type, "L", ll);
             *position += 8;
+            Py_DECREF(bsonint64_type);
             break;
         }
     case 255:
diff --git a/bson/bsonint64.py b/bson/bsonint64.py
@@ -0,0 +1,18 @@
+from bson.py3compat import PY3
+
+if PY3:
+    long = int
+
+
+class BSONInt64(long):
+    """Representation of the BSON int64 type.
+
+    This is necessary because every integral number is an :class:`int` in
+    Python 3. Small integral numbers are encoded to BSON int32 by default,
+    but BSONInt64 numbers will always be encoded to BSON int64.
+
+    :Parameters:
+      - `value`: the numeric value to represent
+    """
+
+    _type_marker = 18
diff --git a/test/test_bson.py b/test/test_bson.py
@@ -30,6 +30,7 @@
                   is_valid,
                   Regex)
 from bson.binary import Binary, UUIDLegacy
+from bson.bsonint64 import BSONInt64
 from bson.code import Code
 from bson.objectid import ObjectId
 from bson.dbref import DBRef
@@ -234,10 +235,10 @@ def helper(dict):
         helper({})
         helper({"test": u("hello")})
         self.assertTrue(isinstance(BSON.encode({"hello": "world"})
-                                .decode()["hello"],
+                                   .decode()["hello"],
                                    text_type))
         helper({"mike": -10120})
-        helper({"long": long(10)})
+        helper({"long": BSONInt64(10)})
         helper({"really big long": 2147483648})
         helper({u("hello"): 0.0013109})
         helper({"something": True})
@@ -372,18 +373,16 @@ def test_overflow(self):
                           {"x": long(-9223372036854775809)})
 
     def test_small_long_encode_decode(self):
-        if PY3:
-            raise SkipTest("No long type in Python 3.")
-
         encoded1 = BSON.encode({'x': 256})
         decoded1 = BSON.decode(encoded1)['x']
         self.assertEqual(256, decoded1)
         self.assertEqual(type(256), type(decoded1))
 
-        encoded2 = BSON.encode({'x': long(256)})
+        encoded2 = BSON.encode({'x': BSONInt64(256)})
         decoded2 = BSON.decode(encoded2)['x']
-        self.assertEqual(long(256), decoded2)
-        self.assertEqual(type(long(256)), type(decoded2))
+        expected = BSONInt64(256)
+        self.assertEqual(expected, decoded2)
+        self.assertEqual(type(expected), type(decoded2))
 
         self.assertNotEqual(type(decoded1), type(decoded2))
 
diff --git a/test/test_database.py b/test/test_database.py
@@ -21,6 +21,7 @@
 
 sys.path[0:0] = [""]
 
+from bson.bsonint64 import BSONInt64
 from bson.code import Code
 from bson.regex import Regex
 from bson.dbref import DBRef
@@ -742,7 +743,14 @@ def test_long(self):
         db = self.client.pymongo_test
         db.test.remove({})
         db.test.save({"x": long(9223372036854775807)})
-        self.assertEqual(long(9223372036854775807), db.test.find_one()["x"])
+        retrieved = db.test.find_one()['x']
+        self.assertEqual(BSONInt64(9223372036854775807), retrieved)
+        self.assertIsInstance(retrieved, BSONInt64)
+        db.test.remove({})
+        db.test.insert({"x": BSONInt64(1)})
+        retrieved = db.test.find_one()['x']
+        self.assertEqual(BSONInt64(1), retrieved)
+        self.assertIsInstance(retrieved, BSONInt64)
 
     def test_remove(self):
         db = self.client.pymongo_test
