
                         p    S r SSKJr  SSKJr  SSKJr  SSKrSSKJr  SSKJr  SSK	J
r
  S	r " S
 S5      rg)z.Utilities for representing a part of a stream.    )absolute_import)division)unicode_literalsN)errors)	hash_util)
installersi   c                       \ rS rSrSr   SS jrS rS rSS jrS r	S	 r
S
 rSS jr\R                  4S jrSS jrS rS rSrg)UploadStream   zImplements a subset of the io.IOBase API, adding functionality for uploads.

When data is read from a stream, this class
1. Updates hash digesters.
2. Executes a progress callbacks if a byte threshold is passed.
Nc                     Xl         X l        Ub  UO0 U l        X@l        SU l        SU l        SU l        SU l        SU l        g)a  Initializes a FilePart instance.

Args:
  stream (io.IOBase): The underlying stream wrapped by this class.
  length (int|None): The total number of bytes in the UploadStream.
  digesters (dict[util.HashAlgorithm, hashlib hash object]|None): Values are
    updated with with data as it's read.
  progress_callback (func[int]|None): Accepts an amount of processed bytes
    and submits progress information for aggregation.
Nr   F)	_stream_length
_digesters_progress_callback(_bytes_read_since_last_progress_callback_progress_updated_with_end_byte_checkpoint_digesters_checkpoint_absolute_index_start_byte)selfstreamlength	digestersprogress_callbacks        7lib/googlecloudsdk/command_lib/storage/upload_stream.py__init__UploadStream.__init__&   sJ     LL#,#8ibDO/45D1+0D(!%D&'D#D    c                 6    U R                   R                  5       $ )a:  Returns absolute position in the stream.

Hashing and progress reporting logic relies on absolute positions. Since
child classes overwrite `tell` to make it return relative positions, we need
to write hashing and progress reporting in a way that does not reference
`self.tell`, which this function makes possible.
)r   tellr   s    r   _get_absolute_position#UploadStream._get_absolute_positionA   s     <<r   c                 8    U R                   R                  U5      $ )a  Seeks to a position in the underlying stream.

Catching up digesters sometimes requires seeking to a specific position in
self._stream. Child classes wrap streams which are not seekable, and have
different strategies to make it appear that a seek has occured, which can
be supported by overriding this method.

Args:
  offset (int): the position to seek to.

Returns:
  the new position in the stream.
)r   seek)r   offsets     r   _update_absolute_position&UploadStream._update_absolute_positionK   s     <<V$$r   c                 8    U R                   R                  U5      $ )a<  Reads bytes from the underlying stream.

Child classes do not always read directly from the stream. Progress
reporting and hashing logic can be reused by overriding only this method.

Args:
  size (int): the number of bytes to read. If less than 0, all bytes are
      returned.

Returns:
  bytes from self._stream.
)r   read)r   sizes     r   	_get_dataUploadStream._get_data[   s     <<T""r   c                     U R                   (       d  gU R                  5       U l        [        R                  " U R                   5      U l        g)zEUpdates checkpoint that holds old hashes to optimize backwards seeks.N)r   r"   r   r   copy_digestersr   r!   s    r   _save_digesters_checkpoint'UploadStream._save_digesters_checkpointj   s4    ??&*&A&A&CD#!*!9!9$//!JDr   c                    U R                   (       d  gXR                  :  a<  U R                  U R                  5        [        R
                  " U R                   5        OXR                  5       :  af  U R                  U R                  5        U R                   R                  U R                  5        [        R                  " U R                  5      U l        OXR                  5       :X  a  gU R                  5          U R                  [        XR                  5       -
  [        R                  5      5      nU(       d  g[        R                  " U R                   U5        Md  )z6Digests data between last and current stream position.N)r   r   r'   r   r   reset_digestersr"   updater   r/   r0   r,   minr   WRITE_BUFFER_SIZEupdate_digesters)r   new_absolute_indexdatas      r   _catch_up_digesters UploadStream._catch_up_digestersq   s	   ??;;;
$$T%5%560	99;	;
$$T%D%DE ooT778#,#;#;

$
$$&d 	::<	< 	##%
^^
 #>#>#@@**,-d   $7 r   c                 "    U R                  5       $ )z+Returns the current position in the stream.)r"   r!   s    r   r    UploadStream.tell   s    &&((r   c                    U R                  5         U R                  U5      nU(       a  [        R                  " U R                  U5        U R
                  (       az  U =R                  [        U5      -  sl        U R                  [        :  aH  SU l        U R                  U R                  5       5        U R                  5       U R                  :H  U l        U$ )z0Returns `size` bytes from the underlying stream.r   )r0   r,   r   r7   r   r   r   len_PROGRESS_CALLBACK_THRESHOLDr"   r    r   r   )r   r+   r9   s      r   r*   UploadStream.read   s    ##%>>$D  $7		 	 55TB599():;$
7

!
!$"="="?
@151L$
.Kr   c                 .   U[         R                  :X  a6  U R                  (       a  XR                  -   nO@[        R                  " S5      eU[         R
                  :X  a  U R                  5       U-   nOUnU R                  U5        U R                  U5      $ )a)  Goes to a specific point in the stream.

Args:
  offset (int): The number of bytes to move.
  whence: Specifies the position offset is added to.
    os.SEEK_SET: offset is added to the current byte.
    os.SEEK_END, os.SEEK_CUR are not supported.

Returns:
  The new position in the stream (int).
zASEEK_END is not supported if the length of the stream is unknown.)	osSEEK_ENDr   r   ErrorSEEK_CURr"   r:   r'   )r   r&   whencer8   s       r   r%   UploadStream.seek   s     	#ll2llOQ 	Q	2;;	6686A!/0))*<==r   c                     U R                   (       a>  U R                  (       d-  U(       d  U R                  U R                  5       5        SU l        U R                  R	                  5       $ )zCloses the underlying stream.T)r   r   r"   r   close)r   caught_errors     r   rJ   UploadStream.close   sH    (L(L ; ; =>-1d*<<r   c                     U $ )N r!   s    r   	__enter__UploadStream.__enter__   s    Kr   c                 4    U R                  [        U5      S9  g )N)rK   )rJ   bool)r   
error_typeunused_argss      r   __exit__UploadStream.__exit__   s    JJD,J-r   )	r   r   r   r   r   r   r   r   r   )NNN))F)__name__
__module____qualname____firstlineno____doc__r   r"   r'   r,   r0   r:   r    r*   rC   SEEK_SETr%   rJ   rO   rU   __static_attributes__rN   r   r   r
   r
      sY     !%	6% #K 8D)  !# >6 .r   r
   )r\   
__future__r   r   r   rC   "googlecloudsdk.command_lib.storager   r   googlecloudsdk.core.updaterr   r@   r
   rN   r   r   <module>rb      s1    5 &  ' 	 5 8 2' p. p.r   