feat(firestore): add zero-dependency PyMongo duck-typing write support (PR 1B)

Author: ohmayr

packages/google-cloud-firestore/google/cloud/firestore_v1/_helpers.py

@@ -44,6 +44,12 @@
 import google
 from google.cloud import exceptions  # type: ignore
 from google.cloud.firestore_v1 import transforms, types
+from google.cloud.firestore_v1.bson import (
+    BSONBinary,
+    BSONDecimal128,
+    BSONObjectID,
+    BSONRegex,
+)
 from google.cloud.firestore_v1.field_path import FieldPath, parse_field_path
 from google.cloud.firestore_v1.types import common, document, write
 from google.cloud.firestore_v1.types.write import DocumentTransform
@@ -185,6 +191,18 @@ def encode_value(value) -> types.document.Value:
     if hasattr(value, "to_map_value"):
         return encode_value(value.to_map_value())
 
+    # Duck-typing input bridge for external PyMongo / bson package objects (zero dependency)
+    if hasattr(value, "binary") and not isinstance(
+        value, (bytes, bytearray, BSONBinary)
+    ):
+        return encode_value(BSONObjectID(getattr(value, "binary")))
+    elif hasattr(value, "to_decimal") and not isinstance(value, BSONDecimal128):
+        return encode_value(BSONDecimal128(getattr(value, "to_decimal")()))
+    elif hasattr(value, "pattern") and not isinstance(value, (str, BSONRegex)):
+        return encode_value(
+            BSONRegex(getattr(value, "pattern"), str(getattr(value, "flags", "")))
+        )
+
     # Must come before int since ``bool`` is an integer subtype.
     if isinstance(value, bool):
         return document.Value(boolean_value=value)

packages/google-cloud-firestore/tests/system/test_bson.py

@@ -66,3 +66,40 @@ def test_bson_document_writes(client, cleanup, database):
 
     snapshot = doc_ref.get()
     assert snapshot.exists
+
+
+@pytest.mark.parametrize("database", [FIRESTORE_ENTERPRISE_DB], indirect=True)
+def test_pymongo_document_writes(client, cleanup, database):
+    """Test write operations using native duck-typed PyMongo objects on Enterprise DB."""
+    import decimal
+
+    class DummyPyMongoObjectId:
+        def __init__(self, raw: bytes):
+            self.binary = raw
+
+    class DummyPyMongoDecimal128:
+        def __init__(self, d: decimal.Decimal):
+            self._d = d
+
+        def to_decimal(self):
+            return self._d
+
+    class DummyPyMongoRegex:
+        def __init__(self, pat: str, flags: str):
+            self.pattern = pat
+            self.flags = flags
+
+    collection_id = "pymongo_docs_write_" + UNIQUE_RESOURCE_ID
+    doc_ref = client.collection(collection_id).document("pymongo_doc")
+    cleanup(doc_ref.delete)
+
+    payload = {
+        "_id": DummyPyMongoObjectId(bytes.fromhex("507f191e810c19729de860ea")),
+        "price": DummyPyMongoDecimal128(decimal.Decimal("99.99")),
+        "pattern": DummyPyMongoRegex("^test.*", "i"),
+    }
+
+    doc_ref.set(payload)
+
+    snapshot = doc_ref.get()
+    assert snapshot.exists

packages/google-cloud-firestore/tests/unit/v1/test__helpers.py

@@ -354,6 +354,45 @@ def test_encode_dict_w_many_types():
     assert encoded_dict == expected_dict
 
 
+def test_encode_value_duck_typed_pymongo():
+    import decimal
+
+    from google.cloud.firestore_v1._helpers import encode_value
+
+    class DummyPyMongoObjectId:
+        def __init__(self, raw: bytes):
+            self.binary = raw
+
+    class DummyPyMongoDecimal128:
+        def __init__(self, d: decimal.Decimal):
+            self._d = d
+
+        def to_decimal(self):
+            return self._d
+
+    class DummyPyMongoRegex:
+        def __init__(self, pat: str, flags: str):
+            self.pattern = pat
+            self.flags = flags
+
+    dummy_oid = DummyPyMongoObjectId(bytes.fromhex("507f1f77bcf86cd799439011"))
+    res_oid = encode_value(dummy_oid)
+    assert (
+        res_oid.map_value.fields["__oid__"].string_value == "507f1f77bcf86cd799439011"
+    )
+
+    dummy_dec = DummyPyMongoDecimal128(decimal.Decimal("99.99"))
+    res_dec = encode_value(dummy_dec)
+    assert res_dec.map_value.fields["__decimal128__"].string_value == "99.99"
+
+    dummy_reg = DummyPyMongoRegex("^test$", "i")
+    res_reg = encode_value(dummy_reg)
+    assert (
+        res_reg.map_value.fields["__regex__"].map_value.fields["pattern"].string_value
+        == "^test$"
+    )
+
+
 def test_reference_value_to_document_w_bad_format():
     from google.cloud.firestore_v1._helpers import (
         BAD_REFERENCE_ERROR,