Merge pull request #40 from davidlatwe/fix-#39

Fix #39, support insert bytes without bson
davidlatwe · Jun 7, 2021 · b35b27c · b35b27c
2 parents 556cd2a + d477e3e
commit b35b27c
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 0 deletions.
diff --git a/montydb/types/_bson.py b/montydb/types/_bson.py
@@ -1,5 +1,6 @@
 import sys
 import types
+import base64
 
 
 class BSON_(types.ModuleType):
@@ -204,6 +205,10 @@ def default(self, obj):
             if isinstance(obj, NoBSON.ObjectId):
                 return {"$oid": str(obj)}
 
+            if isinstance(obj, bytes):
+                return {"$binary": {"base64": base64.b64encode(obj).decode(),
+                                    "subType": "00"}}
+
             if isinstance(obj, NoBSON.datetime.datetime):
                 millis = NoBSON._datetime_to_millis(obj)
                 return {"$date": millis}
@@ -252,6 +257,8 @@ def document_encode(cls, doc, check_keys=False, *args, **kwargs):
     def object_hook(cls, obj, opts=DEFAULT_CODEC_OPTIONS):
         if "$oid" in obj:
             return cls.ObjectId(obj["$oid"])
+        if "$binary" in obj:
+            return base64.b64decode(obj["$binary"]["base64"])
         if "$date" in obj:
             return cls._millis_to_datetime(int(obj["$date"]), opts)
         if "$regex" in obj:

diff --git a/tests/test_engine/test_insert.py b/tests/test_engine/test_insert.py
@@ -30,3 +30,15 @@ def test_insert_invalid_doc_2(monty_find, mongo_find):
 
     with pytest.raises(errors.InvalidDocument):
         monty_find(docs, spec)
+
+
+def test_insert_bytes(monty_find, mongo_find):
+    docs = [
+        {"data": b"some bytes"},
+    ]
+    spec = {}
+
+    monty_c = monty_find(docs, spec)
+    mongo_c = mongo_find(docs, spec)
+
+    assert next(monty_c) == next(mongo_c)