From 7cfced925e8b90d203d0cf2182e107dc0da500ae Mon Sep 17 00:00:00 2001
From: Steve Dignam <steve@dignam.xyz>
Date: Thu, 28 Jan 2021 17:13:49 -0500
Subject: [PATCH] fix(bson): missing top level BSON class (#6)

---
 s/lint                    |  2 +-
 tests/test_mongoengine.py |  8 +++++++-
 typings/bson/__init__.pyi | 15 ++++++++++++++-
 3 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/s/lint b/s/lint
index 0aa5fd1..e2f68d7 100755
--- a/s/lint
+++ b/s/lint
@@ -1,6 +1,6 @@
 #!/bin/bash
 
-set -eux
+set -ex
 
 # format code
 if [[ $CI ]]; then
diff --git a/tests/test_mongoengine.py b/tests/test_mongoengine.py
index 9c93a10..f84badb 100644
--- a/tests/test_mongoengine.py
+++ b/tests/test_mongoengine.py
@@ -6,7 +6,7 @@
 
 import mongoengine
 import pymongo
-from bson import ObjectId
+from bson import BSON, ObjectId
 from mongoengine import Document, EmbeddedDocument, QuerySet, fields
 from pymongo.collation import Collation, CollationStrength
 
@@ -173,6 +173,12 @@ def main() -> None:
     Post().reload()
 
 
+def test_bson() -> None:
+    doc = {"foo": "bar", "buzz": True}
+    data = BSON.encode(doc)
+    assert BSON.decode(data) == doc
+
+
 def test_pymongo() -> None:
     pymongo.ALL
     pymongo.ASCENDING
diff --git a/typings/bson/__init__.pyi b/typings/bson/__init__.pyi
index 3ca609c..db672ce 100644
--- a/typings/bson/__init__.pyi
+++ b/typings/bson/__init__.pyi
@@ -1,4 +1,4 @@
-from typing import Any
+from typing import Any, Mapping
 
 from bson.codec_options import CodecOptions
 from bson.objectid import ObjectId
@@ -6,4 +6,17 @@ from bson.son import SON
 
 def decode_iter(data: bytes, codec_options: CodecOptions = ...) -> Any: ...
 
+class BSON(bytes):
+    @classmethod
+    def encode(
+        cls,
+        document: Mapping[str, Any],
+        check_keys: bool = ...,
+        codec_options: CodecOptions = ...,
+    ) -> BSON: ...
+    def decode(  # type: ignore [override]
+        self,
+        codec_options: CodecOptions = ...,
+    ) -> Any: ...
+
 __all__ = ["ObjectId", "SON", "decode_iter", "CodecOptions"]