4444import google
4545from google .cloud import exceptions # type: ignore
4646from google .cloud .firestore_v1 import transforms , types
47- from google .cloud .firestore_v1 .bson import _BSONType
47+ from google .cloud .firestore_v1 .bson import _BSON_DECODERS , _BSONType
4848from google .cloud .firestore_v1 .field_path import FieldPath , parse_field_path
4949from google .cloud .firestore_v1 .types import common , document , write
5050from google .cloud .firestore_v1 .types .write import DocumentTransform
@@ -347,27 +347,18 @@ def reference_value_to_document(reference_value, client) -> Any:
347347 return document
348348
349349
350- def decode_value (
351- value , client
352- ) -> Union [
353- None , bool , int , float , list , datetime .datetime , str , bytes , dict , GeoPoint , Vector
354- ]:
350+ def decode_value (value , client = None , decode_bson : Optional [bool ] = None ) -> Any :
355351 """Converts a Firestore protobuf ``Value`` to a native Python value.
356352
357353 Args:
358354 value (google.cloud.firestore_v1.types.Value): A
359355 Firestore protobuf to be decoded / parsed / converted.
360356 client (:class:`~google.cloud.firestore_v1.client.Client`):
361357 A client that has a document factory.
358+ decode_bson (Optional[bool]): Whether to decode BSON extended types.
362359
363360 Returns:
364- Union[NoneType, bool, int, float, datetime.datetime, \
365- str, bytes, dict, ~google.cloud.Firestore.GeoPoint]: A native
366- Python value converted from the ``value``.
367-
368- Raises:
369- NotImplementedError: If the ``value_type`` is ``reference_value``.
370- ValueError: If the ``value_type`` is unknown.
361+ Any: A native Python value converted from the ``value``.
371362 """
372363 value_pb = getattr (value , "_pb" , value )
373364 value_type = value_pb .WhichOneof ("value_type" )
@@ -394,37 +385,81 @@ def decode_value(
394385 )
395386 elif value_type == "array_value" :
396387 return [
397- decode_value (element , client ) for element in value_pb .array_value .values
388+ decode_value (element , client , decode_bson = decode_bson )
389+ for element in value_pb .array_value .values
398390 ]
399391 elif value_type == "map_value" :
400- return decode_dict (value_pb .map_value .fields , client )
392+ return decode_dict (value_pb .map_value .fields , client , decode_bson = decode_bson )
401393 else :
402394 raise ValueError ("Unknown ``value_type``" , value_type )
403395
404396
405- def decode_dict (value_fields , client ) -> Union [dict , Vector ]:
397+ def _decode_bson_dict (data : dict ) -> Optional [_BSONType ]:
398+ """Decode a single-key wire map dictionary if registered."""
399+ if len (data ) == 1 :
400+ key , val = next (iter (data .items ()))
401+ decoder = _BSON_DECODERS .get (key )
402+ if decoder is not None :
403+ try :
404+ return decoder (val )
405+ except Exception :
406+ pass
407+ return None
408+
409+
410+ def _decode_bson_dict_recursive (data : Any ) -> Any :
411+ """Recursively decodes BSON wire map dictionaries."""
412+ if isinstance (data , dict ):
413+ decoded = _decode_bson_dict (data )
414+ if decoded is not None :
415+ return decoded
416+ return {k : _decode_bson_dict_recursive (v ) for k , v in data .items ()}
417+ elif isinstance (data , list ):
418+ return [_decode_bson_dict_recursive (item ) for item in data ]
419+ return data
420+
421+
422+ def decode_dict (
423+ value_fields ,
424+ client = None ,
425+ decode_bson : Optional [bool ] = None ,
426+ ) -> Union [dict , Vector , _BSONType ]:
406427 """Converts a protobuf map of Firestore ``Value``-s.
407428
408429 Args:
409430 value_fields (google.protobuf.pyext._message.MessageMapContainer): A
410431 protobuf map of Firestore ``Value``-s.
411432 client (:class:`~google.cloud.firestore_v1.client.Client`):
412433 A client that has a document factory.
434+ decode_bson (Optional[bool]): Whether to decode BSON extended types.
413435
414436 Returns:
415- Dict[str, Union[NoneType, bool, int, float, datetime.datetime , \
416- str, bytes, dict, ~google.cloud.Firestore.GeoPoint]] : A dictionary
417- of native Python values converted from the ``value_fields``.
437+ Union[dict, ~google.cloud.firestore_v1.vector.Vector , \
438+ ~google.cloud.firestore_v1.bson._BSONType] : A dictionary of native \
439+ Python values, Vector, or BSON object converted from ``value_fields``.
418440 """
419441 value_fields_pb = getattr (value_fields , "_pb" , value_fields )
420- res = {key : decode_value (value , client ) for key , value in value_fields_pb .items ()}
442+ res = {
443+ key : decode_value (value , client , decode_bson = decode_bson )
444+ for key , value in value_fields_pb .items ()
445+ }
421446
422447 if res .get ("__type__" , None ) == "__vector__" :
423448 # Vector data type is represented as mapping.
424449 # {"__type__":"__vector__", "value": [1.0, 2.0, 3.0]}.
425450 values = cast (Sequence [float ], res ["value" ])
426451 return Vector (values )
427452
453+ should_decode = (
454+ decode_bson
455+ if decode_bson is not None
456+ else getattr (client , "_decode_bson" , False )
457+ )
458+ if should_decode :
459+ decoded = _decode_bson_dict (res )
460+ if decoded is not None :
461+ return decoded
462+
428463 return res
429464
430465
0 commit comments