U
    >i                     @  s  U d Z ddlmZ ddlZddlZddlmZmZmZm	Z	m
Z
mZmZmZmZmZmZ ddlZddlmZ ddlmZ ddlmZ ddlmZ dd	lmZ dd
lmZ ddlmZ ddl Z ddl!m"Z" ddl#m$Z$m%Z% ddl&m'Z'm(Z( ddl)m*Z*m+Z+m,Z, ddl-m.Z. ddl/m0Z0 er$ddl#m1Z1 de2d< de2d< dZ3dZ4dZ5dZ6dZ7dZ8e.j9j:j;Z<ej=j>e"j?ej=j@e"jAiZBG dd deCZDdd d!d"ZEd#d d$d%ZFdd d&d'ZGd(d)d*d+d,ZHG d-d. d.ZId/d d0d1ZJd2d d3d4ZKd5d d6d7ZLd8d d9d:ZMe$Nd;ZOdd=d>d?d@dAZPdd dBdCZQd/d dDdEZRG dFdG dGeCZSdHd dIdJZTdHd dKdLZUG dMdN dNeSZVdHd dOdPZWG dQdR dReSZXdHd dSdTZYdUd dVdWZZG dXdY dYe[Z\dd[d d\d]Z]d8d^d_d`daZ^G dbdc dceCZ_G ddde dee_Z`G dfdg dge_ZadhdiddjdkdlZbddmdndodpdqZcdrdrdsdtdudvZddwdxdydzd{d|Zedwd}d~ddZfdddddddZgd8ddddZhdS )z<Common helpers shared across Google Cloud Firestore modules.    )annotationsN)AnyDict	GeneratorIteratorListOptionalSequenceTupleUnioncastTYPE_CHECKING)gapic_v1)retry)DatetimeWithNanoseconds)_datetime_to_pb_timestamp)
struct_pb2)	Timestamp)
latlng_pb2)
exceptions)
transformstypes)	FieldPathparse_field_path)commondocumentwrite)DocumentTransform)Vector)DocumentSnapshotztransforms.Sentinel
_EmptyDictdict_GRPC_ERROR_MAPPINGz<A path element must be a string. Received {}, which is a {}./z<Transaction not in progress, cannot be used in API requests.z,Attempted read after write in a transaction.zReference value {!r} in unexpected format, expected to be of the form ``projects/{{project}}/databases/{{database}}/documents/{{document_path}}``.zLDocument {!r} does not correspond to the same database ({!r}) as the client.c                   @  s<   e Zd ZdZddddZddddZd	d
 Zdd ZdS )GeoPointzSimple container for a geo point value.

    Args:
        latitude (float): Latitude of a point.
        longitude (float): Longitude of a point.
    Nonereturnc                 C  s   || _ || _d S Nlatitude	longitude)selfr*   r+    r-   F/tmp/pip-unpacked-wheel-599y42ri/google/cloud/firestore_v1/_helpers.py__init__T   s    zGeoPoint.__init__zlatlng_pb2.LatLngc                 C  s   t j| j| jdS )zConvert the current object to protobuf.

        Returns:
            google.type.latlng_pb2.LatLng: The current point as a protobuf.
        r)   )r   ZLatLngr*   r+   r,   r-   r-   r.   to_protobufX   s    zGeoPoint.to_protobufc                 C  s&   t |tstS | j|jko$| j|jkS )zCompare two geo points for equality.

        Returns:
            Union[bool, NotImplemented]: :data:`True` if the points compare
            equal, else :data:`False`. (Or :data:`NotImplemented` if
            ``other`` is not a geo point.)
        )
isinstancer$   NotImplementedr*   r+   r,   otherr-   r-   r.   __eq__`   s    
zGeoPoint.__eq__c                 C  s    |  |}|tkrtS | S dS )zCompare two geo points for inequality.

        Returns:
            Union[bool, NotImplemented]: :data:`False` if the points compare
            equal, else :data:`True`. (Or :data:`NotImplemented` if
            ``other`` is not a geo point.)
        N)r6   r3   )r,   r5   Zequality_valr-   r-   r.   __ne__m   s    
zGeoPoint.__ne__N)__name__
__module____qualname____doc__r/   r1   r6   r7   r-   r-   r-   r.   r$   L   s
   r$   r%   r&   c                 C  sv   t | }|dkrtd|r2|d dkrFtdn|d dkrFtd| D ]&}t|tsJt|t|}t|qJdS )aR  Verifies that a ``path`` has the correct form.

    Checks that all of the elements in ``path`` are strings.

    Args:
        path (Tuple[str, ...]): The components in a collection or
            document path.
        is_collection (bool): Indicates if the ``path`` represents
            a document or a collection.

    Raises:
        ValueError: if

            * the ``path`` is empty
            * ``is_collection=True`` and there are an even number of elements
            * ``is_collection=False`` and there are an odd number of elements
            * an element is not a string
    r   z+Document or collection path cannot be empty   z5A collection must have an odd number of path elements   z4A document must have an even number of path elementsN)len
ValueErrorr2   strBAD_PATH_TEMPLATEformattype)pathZis_collectionZnum_elementselementmsgr-   r-   r.   verify_path|   s    

rG   ztypes.document.Valuec                 C  s  | dkrt jtjdS t| tr,t j| dS t| trBt j| dS t| trXt j| dS t| trrt j| 	 dS t| t
j
rt jt| dS t| trt j| dS t| trt j| dS t| d	d}|dk	rt j|d
S t| trt j|  dS t| ttttfr2tdd | D }t j|d}t j|dS t| trJt|  S t| trvt| }t j|d}t j|dS td| dt| dS )a  Converts a native Python value into a Firestore protobuf ``Value``.

    Args:
        value (Union[NoneType, bool, int, float, datetime.datetime,             str, bytes, dict, ~google.cloud.Firestore.GeoPoint,             ~google.cloud.firestore_v1.vector.Vector]): A native
            Python value to convert to a protobuf field.

    Returns:
        ~google.cloud.firestore_v1.types.Value: A
        value encoded as a Firestore protobuf.

    Raises:
        TypeError: If the ``value`` is not one of the accepted types.
    N)
null_value)boolean_value)integer_value)double_value)timestamp_value)string_value)bytes_value_document_pathreference_value)geo_point_valuec                 s  s   | ]}t |V  qd S r(   encode_value.0rE   r-   r-   r.   	<genexpr>   s     zencode_value.<locals>.<genexpr>values)array_value)fields)	map_valuez#Cannot convert to a Firestore ValuezInvalid type)r   Valuer   Z
NULL_VALUEr2   boolintfloatr   Ztimestamp_pbdatetimer   r@   bytesgetattrr$   r1   listtupleset	frozenset
ArrayValuer   rT   Zto_map_valuer!   encode_dictZMapValue	TypeErrorrC   )valuedocument_path
value_listvalue_pbZ
value_dictr-   r-   r.   rT      sJ    






   rT   c                 C  s   dd |   D S )a*  Encode a dictionary into protobuf ``Value``-s.

    Args:
        values_dict (dict): The dictionary to encode as protobuf fields.

    Returns:
        Dict[str, ~google.cloud.firestore_v1.types.Value]: A
        dictionary of string keys and ``Value`` protobufs as dictionary
        values.
    c                 S  s   i | ]\}}|t |qS r-   rS   rV   keyrk   r-   r-   r.   
<dictcomp>   s      zencode_dict.<locals>.<dictcomp>)items)Zvalues_dictr-   r-   r.   ri      s    ri   z'DocumentSnapshot'z4Optional['google.cloud.firestore_v1.types.Document'])snapshotr'   c                 C  s4   ddl m} | jsd S || jjt| j| j| jdS )Nr   )Document)namer[   create_timeupdate_time)	google.cloud.firestore_v1.typesrt   exists	referencerO   ri   _datarv   rw   )rs   rt   r-   r-   r.   document_snapshot_to_protobuf   s    r|   c                   @  sB   e Zd ZdZddddZeddddZeddd	d
ZdS )DocumentReferenceValuea  DocumentReference path container with accessors for each relevant chunk.

    Usage:
        doc_ref_val = DocumentReferenceValue(
            'projects/my-proj/databases/(default)/documents/my-col/my-doc',
        )
        assert doc_ref_val.project_name == 'my-proj'
        assert doc_ref_val.collection_name == 'my-col'
        assert doc_ref_val.document_id == 'my-doc'
        assert doc_ref_val.database_name == '(default)'

    Raises:
        ValueError: If the supplied value cannot satisfy a complete path.
    r@   rP   c                 C  sd   || _ |t}t|dk r.t|}t||d | _|d | _|d | _	d
|dd  | _d S )N   r=         r#      )_reference_valuesplitDOCUMENT_PATH_DELIMITERr>   BAD_REFERENCE_ERRORrB   r?   project_namecollection_namedatabase_namejoindocument_id)r,   rQ   partsrF   r-   r-   r.   r/     s    




zDocumentReferenceValue.__init__r&   c                 C  s   d | j| jgS )zSComputed property for a DocumentReference's collection_name and
        document Idr#   )r   r   r   r0   r-   r-   r.   full_key"  s    zDocumentReferenceValue.full_keyc              	   C  s&   | j p$dd| jd| jd| j| jgS )Nr#   ZprojectsZ	databases	documents)r   r   r   r   r   r   r0   r-   r-   r.   	full_path(  s    
z DocumentReferenceValue.full_pathN)r8   r9   r:   r;   r/   propertyr   r   r-   r-   r-   r.   r}     s   r}   r   c                 C  sD   ddl m} t| }||j}|j| kr@t| |j}t	||S )a  Convert a reference value string to a document.

    Args:
        reference_value (str): A document reference value.
        client (:class:`~google.cloud.firestore_v1.client.Client`):
            A client that has a document factory.

    Returns:
        :class:`~google.cloud.firestore_v1.document.DocumentReference`:
            The document corresponding to ``reference_value``.

    Raises:
        ValueError: If the ``reference_value`` is not of the expected
            format: ``projects/{project}/databases/{database}/documents/...``.
        ValueError: If the ``reference_value`` does not come from the same
            project / database combination as the ``client``.
    r   )BaseDocumentReference)
Z'google.cloud.firestore_v1.base_documentr   r}   r   r   rO   WRONG_APP_REFERENCErB   Z_database_stringr?   )rQ   clientr   Zdoc_ref_valuer   rF   r-   r-   r.   reference_value_to_document7  s    
r   zZUnion[None, bool, int, float, list, datetime.datetime, str, bytes, dict, GeoPoint, Vector]c                   s   t | d| }|d}|dkr"dS |dkr0|jS |dkr>|jS |dkrL|jS |dkr`t|jS |d	krn|jS |d
kr||j	S |dkrt
|j S |dkrt|jj|jjS |dkrȇ fdd|jjD S |dkrt|jj S td|dS )a  Converts a Firestore protobuf ``Value`` to a native Python value.

    Args:
        value (google.cloud.firestore_v1.types.Value): A
            Firestore protobuf to be decoded / parsed / converted.
        client (:class:`~google.cloud.firestore_v1.client.Client`):
            A client that has a document factory.

    Returns:
        Union[NoneType, bool, int, float, datetime.datetime,             str, bytes, dict, ~google.cloud.Firestore.GeoPoint]: A native
        Python value converted from the ``value``.

    Raises:
        NotImplementedError: If the ``value_type`` is ``reference_value``.
        ValueError: If the ``value_type`` is unknown.
    _pb
value_typerH   NrI   rJ   rK   rL   rM   rN   rQ   rR   rZ   c                   s   g | ]}t | qS r-   decode_valuerU   r   r-   r.   
<listcomp>  s    z decode_value.<locals>.<listcomp>r\   zUnknown ``value_type``)rc   Z
WhichOneofrI   rJ   rK   r   Zfrom_timestamp_pbrL   rM   rN   r   rQ   r$   rR   r*   r+   rZ   rY   decode_dictr\   r[   r?   )rk   r   rn   r   r-   r   r.   r   U  s<    
 
r   zUnion[dict, Vector]c                   sP   t | d| } fdd| D }|dddkrLttt |d }t|S |S )a   Converts a protobuf map of Firestore ``Value``-s.

    Args:
        value_fields (google.protobuf.pyext._message.MessageMapContainer): A
            protobuf map of Firestore ``Value``-s.
        client (:class:`~google.cloud.firestore_v1.client.Client`):
            A client that has a document factory.

    Returns:
        Dict[str, Union[NoneType, bool, int, float, datetime.datetime,             str, bytes, dict, ~google.cloud.Firestore.GeoPoint]]: A dictionary
        of native Python values converted from the ``value_fields``.
    r   c                   s   i | ]\}}|t | qS r-   r   ro   r   r-   r.   rq     s      zdecode_dict.<locals>.<dictcomp>Z__type__NZ
__vector__rk   )rc   rr   getr   r	   r`   r   )Zvalue_fieldsr   Zvalue_fields_pbresrY   r-   r   r.   r     s    r   r@   c                 C  s.   | j td\}}||kr*td| j d||S )a  Parse a document ID from a document protobuf.

    Args:
        document_pb (google.cloud.firestore_v1.            document.Document): A protobuf for a document that
            was created in a ``CreateDocument`` RPC.
        expected_prefix (str): The expected collection prefix for the
            fully-qualified document name.

    Returns:
        str: The document ID from the protobuf.

    Raises:
        ValueError: If the name does not begin with the prefix.
    r=   zUnexpected document namezExpected to begin with)ru   rsplitr   r?   )Zdocument_pbZexpected_prefixprefixr   r-   r-   r.   
get_doc_id  s    r   zMarker for an empty dict valueFr   %Generator[Tuple[Any, Any], Any, None]prefix_pathr'   c           	      c  s   | s|t fV  nrt|  D ]d\}}|r4t|}nt|}t|j|j  }t|trvt||D ]\}}||fV  q`q||fV  qdS )z7Do depth-first walk of tree, yielding field_path, valueN)	r    sortedrr   r   from_stringr   r2   r!   extract_fields)	document_datar   expand_dotsrp   rk   Zsub_key
field_pathZs_pathZs_valuer-   r-   r.   r     s    
r   c                 C  sB   | }|j dd D ]}||i }q|tkr0i }|||j d < dS )z,Set a value into a document for a field_pathN)r   
setdefaultr    )r   r   rk   currentrE   r-   r-   r.   set_field_value  s    r   c                 C  s<   |j std| }|j d d D ]}|| }q ||j d  S )Nz
Empty pathr   )r   r?   )r   r   r   rE   r-   r-   r.   get_field_value  s    
r   c                   @  s   e Zd ZdZddddZdddd	d
Zedd Zedd ZdddddZ	dddddZ
ddddZdddddZdS )DocumentExtractora  Break document data up into actual data and transforms.

    Handle special values such as ``DELETE_FIELD``, ``SERVER_TIMESTAMP``.

    Args:
        document_data (dict):
            Property names and values to use for sending a change to
            a document.
    r%   r&   c                 C  sL  || _ g | _g | _g | _i | _i | _i | _i | _i | _i | _	d| _
t }| |}|D ]\}}||krv|tkrvd| _
qV|tjkr| j| qV|tjkr| j| qVt|tjr|j| j|< qVt|tjr|j| j|< qVt|tjr|j| j|< qVt|tjr|j| j|< qVt|tjr,|j| j|< qV| j| t| j	|| qVd S )NFT)r   field_pathsdeleted_fieldsserver_timestampsarray_removesarray_unions
incrementsminimumsmaximums
set_fieldsempty_documentr   _get_document_iteratorr    r   ZDELETE_FIELDappendZSERVER_TIMESTAMPr2   ZArrayRemoverY   Z
ArrayUnionZ	Incrementrk   ZMaximumZMinimumr   )r,   r   r   iteratorr   rk   r-   r-   r.   r/     s@    


zDocumentExtractor.__init__r   r   r   c                 C  s   t | j|S r(   r   r   r,   r   r-   r-   r.   r   '  s    z(DocumentExtractor._get_document_iteratorc                 C  s(   t | jp$| jp$| jp$| jp$| jp$| jS r(   )r^   r   r   r   r   r   r   r0   r-   r-   r.   has_transforms,  s    z DocumentExtractor.has_transformsc                 C  s<   t | jt| j t| j t| j t| j t| j S r(   )r   r   rd   r   r   r   r   r   r0   r-   r-   r.   transform_paths7  s    z!DocumentExtractor.transform_pathsF#Optional[types.common.DocumentMask]c                 C  s   d S r(   r-   )r,   allow_empty_maskr-   r-   r.   _get_update_maskB  s    z"DocumentExtractor._get_update_maskNtypes.write.Writec                 C  sD   |d k	rt j|d}nd }tjtj|t| jd| ||d}|S )Nry   )ru   r[   )updateZupdate_maskcurrent_document)	r   Preconditionr   Writer   rt   ri   r   r   )r,   rl   ry   r   r   	update_pbr-   r-   r.   get_update_pbG  s     zDocumentExtractor.get_update_pbz2List[types.write.DocumentTransform.FieldTransform]c                   s   dd  dd | j D  fdd| j D   fdd| j D  dd | j D  dd | j D  d	d | j D  }d
d t|D S )Nc                 S  s   dd | D }t j|dS )Nc                 S  s   g | ]}t |qS r-   rS   rU   r-   r-   r.   r   ]  s     zWDocumentExtractor.get_field_transform_pbs.<locals>.make_array_value.<locals>.<listcomp>rX   )r   rh   )rY   rm   r-   r-   r.   make_array_value\  s    zCDocumentExtractor.get_field_transform_pbs.<locals>.make_array_valuec                 S  s$   g | ]}|t jj| td fqS ))r   Zset_to_server_value)r   r   FieldTransformto_api_reprREQUEST_TIME_ENUMrV   rD   r-   r-   r.   r   a  s   z=DocumentExtractor.get_field_transform_pbs.<locals>.<listcomp>c                   s,   g | ]$\}}|t jj|  |d fqS ))r   Zremove_all_from_arrayr   r   r   r   rV   rD   rY   r   r-   r.   r   k  s   c                   s,   g | ]$\}}|t jj|  |d fqS ))r   Zappend_missing_elementsr   r   r   r-   r.   r   u  s   c                 S  s,   g | ]$\}}|t jj| t|d fqS ))r   	incrementr   r   r   r   rT   rV   rD   rk   r-   r-   r.   r     s    c                 S  s,   g | ]$\}}|t jj| t|d fqS ))r   maximumr   r   r-   r-   r.   r     s    c                 S  s,   g | ]$\}}|t jj| t|d fqS ))r   Zminimumr   r   r-   r-   r.   r     s    c                 S  s   g | ]\}}|qS r-   r-   )rV   rD   	transformr-   r-   r.   r     s     )r   r   rr   r   r   r   r   r   )r,   rl   Zpath_field_transformsr-   r   r.   get_field_transform_pbsY  s4    


'0;z)DocumentExtractor.get_field_transform_pbsc                 C  sD   |  |}tjtj||dd}|d k	r@|jjtj|dj |S )N)r   field_transforms)r   r   )	r   r   r   r   r   r   CopyFromr   r   )r,   rl   ry   r   Ztransform_pbr-   r-   r.   get_transform_pb  s    
 z"DocumentExtractor.get_transform_pb)F)NF)N)r8   r9   r:   r;   r/   r   r   r   r   r   r   r   r   r-   r-   r-   r.   r     s   
-


    Dr   zList[types.write.Write]c                 C  sF   t |}|jrtd|j| dd}|jr@|| }|j| |gS )ai  Make ``Write`` protobufs for ``create()`` methods.

    Args:
        document_path (str): A fully-qualified document path.
        document_data (dict): Property names and values to use for
            creating a document.

    Returns:
        List[google.cloud.firestore_v1.types.Write]: One or two
        ``Write`` protobuf instances for ``create()``.
    z.Cannot apply DELETE_FIELD in a create request.Fr   r   r   r?   r   r   r   update_transformsextend)rl   r   	extractorZ	create_pbfield_transform_pbsr-   r-   r.   pbs_for_create  s    
r   c                 C  sB   t |}|jrtd|| }|jr<|| }|j| |gS )ad  Make ``Write`` protobufs for ``set()`` methods.

    Args:
        document_path (str): A fully-qualified document path.
        document_data (dict): Property names and values to use for
            replacing a document.

    Returns:
        List[google.cloud.firestore_v1.types.Write]: One
        or two ``Write`` protobuf instances for ``set()``.
    zdCannot apply DELETE_FIELD in a set request without specifying 'merge=True' or 'merge=[field_paths]'.r   )rl   r   r   set_pbr   r-   r-   r.   pbs_for_set_no_merge  s    

r   c                      s|   e Zd ZdZdd fddZddddZddd	d
ZddddZddddZddddZ	dddddZ
  ZS )DocumentExtractorForMerge7Break document data up into actual data and transforms.r%   r&   c                   s&   t t| | g | _g | _g | _d S r(   )superr   r/   
data_mergetransform_mergemerge)r,   r   	__class__r-   r.   r/     s    z"DocumentExtractorForMerge.__init__c                 C  s0   t | j| j | _| j| _t | j| j | _d S r(   )r   r   r   r   r   r   r   r0   r-   r-   r.   _apply_merge_all  s    z*DocumentExtractorForMerge._apply_merge_allzGenerator[Any, Any, None]c                 c  s.   |D ]$}t |tr|V  qtt| V  qd S r(   )r2   r   r   )r,   r   Zmerge_fieldr-   r-   r.   _construct_merge_paths  s    
z0DocumentExtractorForMerge._construct_merge_pathsrd   c              	   C  s   t | |}tt|d D ]4}|| ||d   }}||rtd||q|D ]D}|| jkrhqXzt| j	| W qX t
k
r   td|Y qXX qX|S )Nr=   zMerge paths overlap: {}, {}zInvalid merge path: {})r   r   ranger>   eq_or_parentr?   rB   r   r   r   KeyError)r,   r   merge_pathsindexlhsrhs
merge_pathr-   r-   r.   _normalize_merge_paths  s    

z0DocumentExtractorForMerge._normalize_merge_pathsc                   s\  j rtd|}jd d = jd d = |_|D ]<  jkrTj  jD ]} 	|rZj| qZq:i }jD ]}t
j|}t||| q|_fddjD }|rtd|tjj _t jD ]"  fddjD }| qfddjD _fddj D _fd	dj D _d S )
Nz1Cannot merge specific fields with empty document.c                   s   g | ]}| j kr|qS r-   )r   rV   r   r0   r-   r.   r   &  s   
z@DocumentExtractorForMerge._apply_merge_paths.<locals>.<listcomp>z!Cannot delete unmerged fields: {}c                   s   g | ]}  |r|qS r-   )r   )rV   Ztransform_path)r   r-   r.   r   4  s   
c                   s   g | ]}| kr|qS r-   r-   r   merged_transform_pathsr-   r.   r   ;  s     c                   s   i | ]\}}| kr||qS r-   r-   r   r   r-   r.   rq   ?  s    z@DocumentExtractorForMerge._apply_merge_paths.<locals>.<dictcomp>c                   s   i | ]\}}| kr||qS r-   r-   r   r   r-   r.   rq   E  s    )r   r?   r   r   r   r   r   r   r   r   r   r   r   r   r   rB   r   rf   r   r   r   rr   r   )r,   r   r   r   Zmerged_set_fieldsrk   Zunmerged_deleted_fieldsZtranform_merge_pathsr-   )r   r   r,   r.   _apply_merge_paths  sP    










z,DocumentExtractorForMerge._apply_merge_pathsc                 C  s    |dkr|    n
| | d S )NT)r   r   )r,   r   r-   r-   r.   apply_mergeK  s    
z%DocumentExtractorForMerge.apply_mergeFr   c                   s     fdd j D }tj|dS )Nc                   s   g | ]}| j kr| qS r-   )r   r   r   r0   r-   r.   r   U  s   
z>DocumentExtractorForMerge._get_update_mask.<locals>.<listcomp>r   )r   r   DocumentMask)r,   r   
mask_pathsr-   r0   r.   r   Q  s    
z*DocumentExtractorForMerge._get_update_mask)F)r8   r9   r:   r;   r/   r   r   r   r   r   r   __classcell__r-   r-   r   r.   r     s   > r   c                 C  s>   t |}|| || }|jr8|| }|j| |gS )a  Make ``Write`` protobufs for ``set()`` methods.

    Args:
        document_path (str): A fully-qualified document path.
        document_data (dict): Property names and values to use for
            replacing a document.
        merge (Optional[bool] or Optional[List<apispec>]):
            If True, merge all fields; else, merge only the named fields.

    Returns:
        List[google.cloud.firestore_v1.types.Write]: One
        or two ``Write`` protobuf instances for ``set()``.
    )r   r   r   r   r   r   r   )rl   r   r   r   r   r   r-   r-   r.   pbs_for_set_with_merge^  s    


r   c                      sF   e Zd ZdZdd fddZdddd	d
ZdddddZ  ZS )DocumentExtractorForUpdater   r%   r&   c                   s   t t| | tdd |D | _t| j}| jD ]*}| D ]}||kr@td||q@q4| j	D ]}||krftd|qfd S )Nc                 S  s   g | ]}t |qS r-   )r   r   )rV   rp   r-   r-   r.   r     s     z7DocumentExtractorForUpdate.__init__.<locals>.<listcomp>zConflicting field path: {}, {}z"Cannot update with nest delete: {})
r   r  r/   r   top_level_pathsrf   Zlineager?   rB   r   )r,   r   ZtopsZtop_level_pathZancestorr   r   r-   r.   r/   }  s&    

 
z#DocumentExtractorForUpdate.__init__r   r   r   c                 C  s   t | j|ddS )NT)r   r   r   r-   r-   r.   r     s    z1DocumentExtractorForUpdate._get_document_iteratorFztypes.common.DocumentMaskc                 C  s4   g }| j D ]}|| jkr
||  q
tj|dS )Nr   )r  r   r   r   r   r   )r,   r   r   r   r-   r-   r.   r     s
    

z+DocumentExtractorForUpdate._get_update_mask)F)r8   r9   r:   r;   r/   r   r   r   r-   r-   r   r.   r  z  s   r  c                 C  s^   t |}|jrtd|dkr(tdd}|| }|| |jrX|| }|j	| |gS )aC  Make ``Write`` protobufs for ``update()`` methods.

    Args:
        document_path (str): A fully-qualified document path.
        field_updates (dict): Field names or paths to update and values
            to update with.
        option (optional[:class:`~google.cloud.firestore_v1.client.WriteOption`]):
            A write option to make assertions / preconditions on the server
            state of the document before applying changes.

    Returns:
        List[google.cloud.firestore_v1.types.Write]: One
        or two ``Write`` protobuf instances for ``update()``.
    z%Cannot update with an empty document.NTr   )
r  r   r?   ExistsOptionr   modify_writer   r   r   r   )rl   Zfield_updatesoptionr   r   r   r-   r-   r.   pbs_for_update  s    



r  r   c                 C  s"   t j| d}|dk	r|| |S )a  Make a ``Write`` protobuf for ``delete()`` methods.

    Args:
        document_path (str): A fully-qualified document path.
        option (optional[:class:`~google.cloud.firestore_v1.client.WriteOption`]):
            A write option to make assertions / preconditions on the server
            state of the document before applying changes.

    Returns:
        google.cloud.firestore_v1.types.Write: A
        ``Write`` protobuf instance for the ``delete()``.
    )deleteN)r   r   r  )rl   r  Zwrite_pbr-   r-   r.   pb_for_delete  s    
r  c                   @  s   e Zd ZdZdS )ReadAfterWriteErrorzhRaised when a read is attempted after a write.

    Raised by "read" methods that use transactions.
    N)r8   r9   r:   r;   r-   r-   r-   r.   r	    s   r	  TzUnion[bytes, None]c                 C  s>   | dkrdS | j stt|r4t| jdkr4tt| jS dS )a  Get the transaction ID from a ``Transaction`` object.

    Args:
        transaction (Optional[:class:`~google.cloud.firestore_v1.transaction.            Transaction`]):
            An existing transaction that this query will run in.
        read_operation (Optional[bool]): Indicates if the transaction ID
            will be used in a read operation. Defaults to :data:`True`.

    Returns:
        Optional[bytes]: The ID of the transaction, or :data:`None` if the
        ``transaction`` is :data:`None`.

    Raises:
        ValueError: If the ``transaction`` is not in progress (only if
            ``transaction`` is not :data:`None`).
        ReadAfterWriteError: If the ``transaction`` has writes stored on
            it and ``read_operation`` is :data:`True`.
    Nr   )Zin_progressr?   INACTIVE_TXNr>   Z
_write_pbsr	  READ_AFTER_WRITE_ERRORid)ZtransactionZread_operationr-   r-   r.   get_transaction_id  s    r  zList[Tuple[str, str]])r   r'   c                 K  s
   d| fgS )zCreate RPC metadata containing a prefix.

    Args:
        prefix (str): appropriate resource path.

    Returns:
        List[Tuple[str, str]]: RPC metadata with supplied prefix
    zgoogle-cloud-resource-prefixr-   )r   kwr-   r-   r.   metadata_with_prefix  s    	r  c                   @  s    e Zd ZdZdddddZdS )WriteOptionz7Option used to assert a condition on a write operation.Nr%   r&   c                 C  s   t dS )a>  Modify a ``Write`` protobuf based on the state of this write option.

        This is a virtual method intended to be implemented by subclasses.

        Args:
            write (google.cloud.firestore_v1.types.Write): A
                ``Write`` protobuf instance to be modified with a precondition
                determined by the state of this option.
            no_create_msg (Optional[str]): A message to use to indicate that
                a create operation is not allowed.

        Raises:
            NotImplementedError: Always, this method is virtual.
        N)NotImplementedError)r,   r   Zno_create_msgr-   r-   r.   r  	  s    zWriteOption.modify_write)N)r8   r9   r:   r;   r  r-   r-   r-   r.   r    s   r  c                   @  s4   e Zd ZdZddddZdd Zdddd	Zd
S )LastUpdateOptiona*  Option used to assert a "last update" condition on a write operation.

    This will typically be created by
    :meth:`~google.cloud.firestore_v1.client.Client.write_option`.

    Args:
        last_update_time (google.protobuf.timestamp_pb2.Timestamp): A
            timestamp. When set, the target document must exist and have
            been last updated at that time. Protobuf ``update_time`` timestamps
            are typically returned from methods that perform write operations
            as part of a "write result" protobuf or directly.
    r%   r&   c                 C  s
   || _ d S r(   )_last_update_time)r,   Zlast_update_timer-   r-   r.   r/   )  s    zLastUpdateOption.__init__c                 C  s   t || jstS | j|jkS r(   )r2   r   r3   r  r4   r-   r-   r.   r6   ,  s    zLastUpdateOption.__eq__c                 O  s"   t j| jd}|jj|j dS )ac  Modify a ``Write`` protobuf based on the state of this write option.

        The ``last_update_time`` is added to ``write_pb`` as an "update time"
        precondition. When set, the target document must exist and have been
        last updated at that time.

        Args:
            write_pb (google.cloud.firestore_v1.types.Write): A
                ``Write`` protobuf instance to be modified with a precondition
                determined by the state of this option.
            unused_kwargs (Dict[str, Any]): Keyword arguments accepted by
                other subclasses that are unused here.
        )rw   N)r   r   r  r   r   r   r,   r   Zunused_argsZunused_kwargsZcurrent_docr-   r-   r.   r  1  s    zLastUpdateOption.modify_writeNr8   r9   r:   r;   r/   r6   r  r-   r-   r-   r.   r    s   r  c                   @  s4   e Zd ZdZddddZdd Zdddd	Zd
S )r  a  Option used to assert existence on a write operation.

    This will typically be created by
    :meth:`~google.cloud.firestore_v1.client.Client.write_option`.

    Args:
        exists (bool): Indicates if the document being modified
            should already exist.
    r%   r&   c                 C  s
   || _ d S r(   )_exists)r,   ry   r-   r-   r.   r/   N  s    zExistsOption.__init__c                 C  s   t || jstS | j|jkS r(   )r2   r   r3   r  r4   r-   r-   r.   r6   Q  s    zExistsOption.__eq__c                 O  s"   t j| jd}|jj|j dS )aB  Modify a ``Write`` protobuf based on the state of this write option.

        If:

        * ``exists=True``, adds a precondition that requires existence
        * ``exists=False``, adds a precondition that requires non-existence

        Args:
            write (google.cloud.firestore_v1.types.Write): A
                ``Write`` protobuf instance to be modified with a precondition
                determined by the state of this option.
            unused_kwargs (Dict[str, Any]): Keyword arguments accepted by
                other subclasses that are unused here.
        r   N)r   r   r  r   r   r   r  r-   r-   r.   r  V  s    zExistsOption.modify_writeNr  r-   r-   r-   r.   r  C  s   
r  z2retries.Retry | retries.AsyncRetry | object | Nonezfloat | None)r   timeoutr'   c                 C  s,   i }| t jjk	r| |d< |dk	r(||d< |S )zCHelper fo API methods which take optional 'retry' / 'timeout' args.r   Nr  )r   methodDEFAULT)r   r  kwargsr-   r-   r.   make_retry_timeout_kwargsi  s    r  z;Optional[Union[DatetimeWithNanoseconds, datetime.datetime]]r   )dtr'   c                 C  s   t | ptjtjjdS )z7Returns the supplied datetime (or "now") as a Timestamp)tz)r   r   nowra   timezoneutc)r  r-   r-   r.   build_timestampx  s    r!  z#Union[Timestamp, datetime.datetime]r_   )ts1ts2r'   c                 C  sh   t | tst| n| } t |ts(t|n|}| j| jd  }|j|jd  }||krXdS ||krddS dS )Ng    eAr   r=   r   )r2   r   r!  Znanosseconds)r"  r#  Z	ts1_nanosZ	ts2_nanosr-   r-   r.   compare_timestamps  s    r%  zUnion[str, bytes]z-'google.cloud.firestore_v1.client.BaseClient'z/'google.cloud.firestore_bundle.FirestoreBundle')
serializedr   r'   c                 C  sf  ddl m}m} dgdddgdddgdgddgd}|d	 }d
}d
}t| D ]}t| }	t|	dkrptd|	d }
|
|krtd|
 d| z|t	
|}W n< tk
r } z|jd dkrtd|W 5 d
}~X Y nX |d
kr|
dkst|||
 d }|}n|j|||
d ||
 }qLd|kr>tdt||}t||}|j||dd |S )a  Inverse operation to a `FirestoreBundle` instance's `build()` method.

    Args:
        serialized (Union[str, bytes]): The result of `FirestoreBundle.build()`.
            Should be a list of dictionaries in string format.
        client (BaseClient): A connected Client instance.

    Returns:
        FirestoreBundle: A bundle equivalent to that which called `build()` and
            initially created the `serialized` value.

    Raises:
        ValueError: If any of the dictionaries in the list contain any more than
            one top-level key.
        ValueError: If any unexpected BundleElement types are encountered.
        ValueError: If the serialized bundle ends before expected.
    r   )BundleElementFirestoreBundlemetadata
namedQuerydocumentMetadataZ__end__r   )__initial__r)  r*  r+  r   r,  Nr=   z8Expected serialized BundleElement with one top-level keyz"Encountered BundleElement of type z. Expected one of z%'dict' object has no attribute 'find'z[Invalid serialization of datetimes. Cannot deserialize Bundles created from the NodeJS SDK.r  )r   rC   z,Unexpected end to serialized FirestoreBundle)Zgoogle.cloud.firestore_bundler'  r(  _parse_bundle_elements_datard   keysr>   r?   Z	from_jsonjsondumpsAttributeErrorargsAssertionErrorZ_add_bundle_elementr   )r&  r   r'  r(  Zbundle_state_machineZallowed_next_element_typesZmetadata_bundle_elementbundledatar.  rp   Zbundle_elementer-   r-   r.   deserialize_bundle  sV    




r7  zGenerator[Dict, None, None])r&  r'   c                 c  s   t t| tr| n| d}d}t|d}|dkr6dS t|}| rP||7 }q |dkr`tdt|}d}t	|g}d}||k r|
t| |d7 }qzt|dV  q dS )a
  Reads through a serialized FirestoreBundle and yields JSON chunks that
    were created via `BundleElement.to_json(bundle_element)`.

    Serialized FirestoreBundle instances are length-prefixed JSON objects, and
    so are of the form "123{...}57{...}"
    To correctly and safely read a bundle, we must first detect these length
    prefixes, read that many bytes of data, and attempt to JSON-parse that.

    Raises:
        ValueError: If a chunk of JSON ever starts without following a length
            prefix.
    zutf-8 NzExpected length prefixr=   )iterr2   rb   encodenextchr	isnumericr?   r_   	bytearrayr   r/  loadsdecode)r&  Z_serializedZlength_prefixbyte_strZ_length_prefixZ_bytes_counterr-   r-   r.   r-    s(    



r-  )
query_namezOptional[str]z+Generator[('DocumentSnapshot', None, None)])rD  r'   c                c  s:   ddl m} | j D ]}|r,||jjkr,q|jV  qd S )Nr   )_BundledDocument)Z$google.cloud.firestore_bundle.bundlerE  r   rY   r)  Zqueriesrs   )r4  rD  rE  bundled_docr-   r-   r.   _get_documents_from_bundle  s
    rG  zOptional['DocumentSnapshot'])r   r'   c                C  s   | j |}|r|jS d S d S r(   )r   r   rs   )r4  r   rF  r-   r-   r.   _get_document_from_bundle(  s    rH  )F)T)N)ir;   
__future__r   ra   r/  typingr   r   r   r   r   r   r	   r
   r   r   r   ZgrpcZgoogle.api_corer   r   retriesZ google.api_core.datetime_helpersr   Zgoogle.cloud._helpersr   Zgoogle.protobufr   Zgoogle.protobuf.timestamp_pb2r   Zgoogle.typer   ZgoogleZgoogle.cloudr   Zgoogle.cloud.firestore_v1r   r   Z$google.cloud.firestore_v1.field_pathr   r   rx   r   r   r   Z%google.cloud.firestore_v1.types.writer   Z google.cloud.firestore_v1.vectorr   r   __annotations__rA   r   r
  r  r   r   r   ZServerValueZREQUEST_TIMEr   Z
StatusCodeZALREADY_EXISTSConflict	NOT_FOUNDZNotFoundr"   objectr$   rG   rT   ri   r|   r}   r   r   r   r   ZSentinelr    r   r   r   r   r   r   r   r   r  r  r  	Exceptionr	  r  r  r  r  r  r  r!  r%  r7  r-  rG  rH  r-   r-   r-   r.   <module>   s   4
  0%D47
 

 >{&!(& 	b-