Skip to content

Commit 2585841

Browse files
ohmayrohmayr
authored
feat(firestore): add BSONRegex support
1 parent 47b2845 commit 2585841

7 files changed

Lines changed: 178 additions & 0 deletions

File tree

‎.librarian/generator-input/client-post-processing/firestore-integration.yaml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,7 @@ replacements:
7575
BSONMaxKey,
7676
BSONMinKey,
7777
BSONObjectId,
78+
BSONRegex,
7879
BSONTimestamp,
7980
)
8081
from google.cloud.firestore_v1.client import Client
@@ -183,6 +184,7 @@ replacements:
183184
"BSONMaxKey",
184185
"BSONMinKey",
185186
"BSONObjectId",
187+
"BSONRegex",
186188
"BSONTimestamp",
187189
"Client",
188190
"CountAggregation",
@@ -261,6 +263,7 @@ replacements:
261263
BSONMaxKey,
262264
BSONMinKey,
263265
BSONObjectId,
266+
BSONRegex,
264267
BSONTimestamp,
265268
Client,
266269
CollectionGroup,
@@ -324,6 +327,7 @@ replacements:
324327
"BSONMaxKey",
325328
"BSONMinKey",
326329
"BSONObjectId",
330+
"BSONRegex",
327331
"BSONTimestamp",
328332
"Client",
329333
"CountAggregation",

‎packages/google-cloud-firestore/google/cloud/firestore/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
BSONMaxKey,
4141
BSONMinKey,
4242
BSONObjectId,
43+
BSONRegex,
4344
BSONTimestamp,
4445
Client,
4546
CollectionGroup,
@@ -103,6 +104,7 @@
103104
"BSONMaxKey",
104105
"BSONMinKey",
105106
"BSONObjectId",
107+
"BSONRegex",
106108
"BSONTimestamp",
107109
"Client",
108110
"CountAggregation",

‎packages/google-cloud-firestore/google/cloud/firestore_v1/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@
5252
BSONMaxKey,
5353
BSONMinKey,
5454
BSONObjectId,
55+
BSONRegex,
5556
BSONTimestamp,
5657
)
5758
from google.cloud.firestore_v1.client import Client
@@ -160,6 +161,7 @@
160161
"BSONMaxKey",
161162
"BSONMinKey",
162163
"BSONObjectId",
164+
"BSONRegex",
163165
"BSONTimestamp",
164166
"Client",
165167
"CountAggregation",

‎packages/google-cloud-firestore/google/cloud/firestore_v1/bson.py‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
"BSONInt32",
3636
"BSONBinary",
3737
"BSONTimestamp",
38+
"BSONRegex",
3839
]
3940

4041
_OBJECT_ID_BYTES_LEN = 12
@@ -342,3 +343,85 @@ def __eq__(self, other: Any) -> bool:
342343

343344
def __hash__(self) -> int:
344345
return hash((type(self), self._seconds, self._increment))
346+
347+
348+
class BSONRegex(_BSONType):
349+
"""Represents a BSON Regular Expression container for Firestore.
350+
351+
Args:
352+
pattern (str): The regular expression pattern string.
353+
options (Union[str, re.RegexFlag, int], optional): BSON regex option flags
354+
as a string (e.g. "i", "m", "s") or Python `re` flag integer (e.g. `re.I | re.M`).
355+
Defaults to "".
356+
357+
Raises:
358+
TypeError: If pattern is not a string or options is invalid type.
359+
360+
Example:
361+
>>> regex = BSONRegex("^hello.*$", options="i")
362+
>>> regex.pattern
363+
'^hello.*$'
364+
>>> regex.options
365+
'i'
366+
"""
367+
368+
__slots__ = ("_pattern", "_options")
369+
370+
_FLAG_TO_OPTION: Dict[int, str] = {
371+
re.IGNORECASE: "i",
372+
re.LOCALE: "l",
373+
re.MULTILINE: "m",
374+
re.DOTALL: "s",
375+
re.UNICODE: "u",
376+
re.VERBOSE: "x",
377+
}
378+
379+
def __init__(self, pattern: str, options: Union[str, re.RegexFlag, int] = ""):
380+
if not isinstance(pattern, str):
381+
raise TypeError("BSONRegex pattern must be a str.")
382+
383+
if isinstance(options, bool):
384+
raise TypeError("BSONRegex options must be a str or re flag integer.")
385+
386+
if isinstance(options, str):
387+
self._options: str = "".join(sorted(set(options)))
388+
elif isinstance(options, int):
389+
opts = []
390+
for flag, char in self._FLAG_TO_OPTION.items():
391+
if options & flag:
392+
opts.append(char)
393+
self._options = "".join(sorted(opts))
394+
else:
395+
raise TypeError("BSONRegex options must be a str or re flag integer.")
396+
397+
self._pattern: str = pattern
398+
399+
@property
400+
def pattern(self) -> str:
401+
"""str: The regular expression pattern string."""
402+
return self._pattern
403+
404+
@property
405+
def options(self) -> str:
406+
"""str: The normalized BSON regex option flags sorted alphabetically."""
407+
return self._options
408+
409+
def _to_map_value(self) -> Dict[str, Dict[str, str]]:
410+
"""Returns map dictionary representation for wire serialization."""
411+
return {
412+
"__regex__": {
413+
"pattern": self._pattern,
414+
"options": self._options,
415+
}
416+
}
417+
418+
def __repr__(self) -> str:
419+
return f"BSONRegex({self._pattern!r}, options={self._options!r})"
420+
421+
def __eq__(self, other: Any) -> bool:
422+
if isinstance(other, BSONRegex):
423+
return self._pattern == other._pattern and self._options == other._options
424+
return NotImplemented
425+
426+
def __hash__(self) -> int:
427+
return hash((type(self), self._pattern, self._options))

‎packages/google-cloud-firestore/tests/system/test_system.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@
5454
BSONMaxKey,
5555
BSONMinKey,
5656
BSONObjectId,
57+
BSONRegex,
5758
BSONTimestamp,
5859
)
5960
from google.cloud.firestore_v1.vector import Vector
@@ -1296,6 +1297,7 @@ def test_bson_document_writes(client, cleanup, database):
12961297
"int32_val": BSONInt32(42),
12971298
"binary_val_sub128": BSONBinary(b"world", subtype=128),
12981299
"timestamp_val": BSONTimestamp(1700000000, 1),
1300+
"regex_val": BSONRegex("^hello.*$", options="i"),
12991301
}
13001302

13011303
doc_ref.set(bson_payload)
@@ -1314,6 +1316,12 @@ def test_bson_document_writes(client, cleanup, database):
13141316
"increment": 1,
13151317
}
13161318
},
1319+
"regex_val": {
1320+
"__regex__": {
1321+
"pattern": "^hello.*$",
1322+
"options": "i",
1323+
}
1324+
},
13171325
}
13181326

13191327

‎packages/google-cloud-firestore/tests/system/test_system_async.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@
5757
BSONMaxKey,
5858
BSONMinKey,
5959
BSONObjectId,
60+
BSONRegex,
6061
BSONTimestamp,
6162
)
6263
from google.cloud.firestore_v1.query_profile import (
@@ -1269,6 +1270,7 @@ async def test_async_bson_document_writes(client, cleanup, database):
12691270
"int32_val": BSONInt32(42),
12701271
"binary_val_sub128": BSONBinary(b"world", subtype=128),
12711272
"timestamp_val": BSONTimestamp(1700000000, 1),
1273+
"regex_val": BSONRegex("^hello.*$", options="i"),
12721274
}
12731275

12741276
await doc_ref.set(bson_payload)
@@ -1287,6 +1289,12 @@ async def test_async_bson_document_writes(client, cleanup, database):
12871289
"increment": 1,
12881290
}
12891291
},
1292+
"regex_val": {
1293+
"__regex__": {
1294+
"pattern": "^hello.*$",
1295+
"options": "i",
1296+
}
1297+
},
12901298
}
12911299

12921300

‎packages/google-cloud-firestore/tests/unit/v1/test_bson.py‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
import copy
1919
import pickle
20+
import re
2021

2122
import pytest
2223

@@ -26,6 +27,7 @@
2627
BSONMaxKey,
2728
BSONMinKey,
2829
BSONObjectId,
30+
BSONRegex,
2931
BSONTimestamp,
3032
_BSONType,
3133
)
@@ -411,3 +413,72 @@ def test_bson_timestamp_copy():
411413
def test_bson_timestamp_pickle():
412414
ts = BSONTimestamp(100, 1)
413415
assert pickle.loads(pickle.dumps(ts)) == ts
416+
417+
418+
def test_bson_regex_valid():
419+
rx = BSONRegex("^hello.*$", options="i")
420+
assert rx.pattern == "^hello.*$"
421+
assert rx.options == "i"
422+
assert rx._to_map_value() == {
423+
"__regex__": {
424+
"pattern": "^hello.*$",
425+
"options": "i",
426+
}
427+
}
428+
assert repr(rx) == "BSONRegex('^hello.*$', options='i')"
429+
430+
431+
def test_bson_regex_options_sorting_and_deduplication():
432+
rx1 = BSONRegex("foo", options="msi")
433+
assert rx1.options == "ims"
434+
435+
rx2 = BSONRegex("foo", options="mmiis")
436+
assert rx2.options == "ims"
437+
438+
439+
def test_bson_regex_options_from_re_flags():
440+
rx = BSONRegex("foo", options=re.IGNORECASE | re.MULTILINE)
441+
assert rx.options == "im"
442+
443+
444+
@pytest.mark.parametrize(
445+
"pattern_input, options_input, exc_type, match_msg",
446+
[
447+
(123, "i", TypeError, "pattern must be a str"),
448+
(None, "i", TypeError, "pattern must be a str"),
449+
("foo", True, TypeError, "options must be a str or re flag integer"),
450+
("foo", [1, 2], TypeError, "options must be a str or re flag integer"),
451+
],
452+
)
453+
def test_bson_regex_invalid_inputs(pattern_input, options_input, exc_type, match_msg):
454+
with pytest.raises(exc_type, match=match_msg):
455+
BSONRegex(pattern_input, options_input)
456+
457+
458+
def test_bson_regex_equality():
459+
rx1 = BSONRegex("^abc", options="i")
460+
rx2 = BSONRegex("^abc", options="i")
461+
rx3 = BSONRegex("^abc", options="m")
462+
rx4 = BSONRegex("^xyz", options="i")
463+
assert rx1 == rx2
464+
assert rx1 != rx3
465+
assert rx1 != rx4
466+
assert rx1 != "^abc"
467+
468+
469+
def test_bson_regex_hash_and_dict_key():
470+
rx1 = BSONRegex("^abc", options="i")
471+
rx2 = BSONRegex("^abc", options="i")
472+
assert hash(rx1) == hash(rx2)
473+
assert len({rx1, rx2}) == 1
474+
475+
476+
def test_bson_regex_copy():
477+
rx = BSONRegex("^abc", options="i")
478+
assert copy.copy(rx) == rx
479+
assert copy.deepcopy(rx) == rx
480+
481+
482+
def test_bson_regex_pickle():
483+
rx = BSONRegex("^abc", options="i")
484+
assert pickle.loads(pickle.dumps(rx)) == rx

0 commit comments

Comments
 (0)