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

Author: ohmayr

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

@@ -18,6 +18,7 @@
 
 import datetime
 import json
+import re
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -44,6 +45,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
@@ -163,13 +170,28 @@ def verify_path(path, is_collection) -> None:
             raise ValueError(msg)
 
 
-def encode_value(value) -> types.document.Value:
-    """Converts a native Python value into a Firestore protobuf ``Value``.
+_REGEX_FLAG_MAP = (
+    (re.IGNORECASE, "i"),
+    (re.MULTILINE, "m"),
+    (re.DOTALL, "s"),
+    (re.VERBOSE, "x"),
+    (re.LOCALE, "l"),
+)
+
+
+def _extract_regex_options(flags: Union[int, str]) -> str:
+    if isinstance(flags, str):
+        return flags
+    if isinstance(flags, int):
+        return "".join(char for bit, char in _REGEX_FLAG_MAP if flags & bit)
+    return ""
+
+
+def encode_value(value: Any) -> document.Value:
+    """Convert a Python value into a Value protobuf.
 
     Args:
-        value (Union[NoneType, bool, int, float, datetime.datetime, \
-            str, bytes, dict, ~google.cloud.Firestore.GeoPoint, \
-            ~google.cloud.firestore_v1.vector.Vector]): A native
+        value (Any): The
             Python value to convert to a protobuf field.
 
     Returns:
@@ -186,6 +208,23 @@ def encode_value(value) -> types.document.Value:
     if callable(to_map):
         return encode_value(to_map())
 
+    # Duck-typing input bridge for external PyMongo / bson package objects (zero dependency)
+    binary_attr = getattr(value, "binary", None)
+    if binary_attr is not None and not isinstance(
+        value, (bytes, bytearray, BSONBinary)
+    ):
+        return encode_value(BSONObjectID(binary_attr))
+
+    to_decimal_fn = getattr(value, "to_decimal", None)
+    if callable(to_decimal_fn) and not isinstance(value, BSONDecimal128):
+        return encode_value(BSONDecimal128(to_decimal_fn()))
+
+    pattern_attr = getattr(value, "pattern", None)
+    if pattern_attr is not None and not isinstance(value, (str, BSONRegex)):
+        flags_attr = getattr(value, "flags", "")
+        options_str = _extract_regex_options(flags_attr)
+        return encode_value(BSONRegex(pattern_attr, options_str))
+
     # 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,
@@ -2572,3 +2611,17 @@ def _make_field_path(*fields):
     from google.cloud.firestore_v1 import field_path
 
     return field_path.FieldPath(*fields)
+
+
+def test_encode_value_w_compiled_regex_flags():
+    import re
+
+    from google.cloud.firestore_v1._helpers import encode_value
+
+    compiled_re = re.compile("abc", re.I | re.M)
+    encoded = encode_value(compiled_re)
+    # Checks that compiled regex integer flags translate to 'im' options
+    fields = encoded.map_value.fields
+    assert fields["__regex__"].map_value.fields["pattern"].string_value == "abc"
+    assert "i" in fields["__regex__"].map_value.fields["options"].string_value
+    assert "m" in fields["__regex__"].map_value.fields["options"].string_value