o
    F                  
   @   s&  d Z ddlmZ ddlmZ ddlmZ ddlZddlZddlZddlZddl	Z	ddl
Z
ddl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ZeZe	jdkrfejZdddddddejjdf	ddZG dd deejeZG dd deZ G dd deZ!G dd deZ"G dd de"Z#dd Z$		d9d!d"Z%G d#d$ d$eZ&G d%d& d&eZ'ddddddddejjdf
d'd(Z(G d)d* d*eZ)G d+d, d,ej*Z+G d-d. d.ej,Z-G d/d0 d0e-Z.G d1d2 d2e.Z/G d3d4 d4e-Z0G d5d6 d6e-Z1G d7d8 d8e1Z2dS ):zProgress Tracker for Cloud SDK.    )absolute_import)division)unicode_literalsN)log)
properties)console_attr)
console_io)	multiline)parser)      Tg?Fc	              
   C   s   t jjj }	|	t jjjjjkrt||S |	t jjjj	jkr$t
| ||S tjdd}
|
r.tnt}|p8t jjj }d}|rAd}d}|| ||||||||	S )a  A context manager for telling the user about long-running progress.

  Args:
    message: str, The message to show next to the spinner.
    autotick: bool, True to have the spinner tick on its own. Otherwise, you
      need to call Tick() explicitly to move the spinner.
    detail_message_callback: func, A no argument function that will be called
      and the result appended to message each time it needs to be printed.
    done_message_callback: func, A no argument function whose result will be
      appended to message if the progress tracker successfully exits.
    tick_delay: float, The amount of time to wait between ticks, in second.
    interruptable: boolean, True if the user can ctrl-c the operation. If so,
      it will stop and will report as aborted. If False, a message will be
      displayed saying that it cannot be cancelled.
    screen_reader: boolean, override for screen reader accessibility property
      toggle.
    aborted_message: str, A custom message to put in the exception when it is
      cancelled by the user.
    no_spacing: boolean, Removes ellipses and other spacing between text.

  Returns:
    The progress tracker.
  TerrorN   working)r   VALUEScoreinteractive_ux_styleGetInteractiveUXStylesOFFnameNoOpProgressTrackerTESTING_StubProgressTrackerr   IsInteractive_NormalProgressTracker_NonInteractiveProgressTrackeraccessibilityscreen_readerGetBool)messageautotickdetail_message_callbackdone_message_callback
tick_delayinterruptabler   aborted_message
no_spacingstyleis_ttytracker_clsspinner_override_message r-   I/tmp/google-cloud-sdk/lib/googlecloudsdk/core/console/progress_tracker.pyProgressTracker-   s(   !
r/   c                   @   sj   e Zd ZdZdd Zedd Zdd Zdd	 Zd
d Z	dd Z
dd Zejdd ZejdddZdS )_BaseProgressTrackerCA context manager for telling the user about long-running progress.c
                 C   s   t j| _|d u rd| _d| _d| _nd| _|| _||	rdnd | _|| _|| _|| _d| _	d| _
t | _|| _d | _t  d }
|
dk rId}
t oP|
dk| _|oV| j| _|| _|| _d | _t  | _|	| _tjdd| _ d S )NT F...r   r   )!sysstderr_stream_spinner_only_message_prefix_detail_message_callbackr,   _done_message_callback_ticks_done	threadingLock_lock_tick_delay_tickerr   ConsoleAttrGetTermSizer   IsUserOutputEnabled_output_enabled_BaseProgressTracker__autotick_interruptable_aborted_message_old_signal_handlerGetConsoleAttrGetProgressTrackerSymbols_symbols_no_spacingr   r   _is_tty)selfr!   r"   r#   r$   r%   r&   r'   r,   r(   console_widthr-   r-   r.   __init__f   s6   
z_BaseProgressTracker.__init__c                 C      | j S N)rG   rP   r-   r-   r.   	_autotick      z_BaseProgressTracker._autotickc                 C   s@   | j r| jr|  }|r| jr| j| S | jd | d S | jS N r3   )rO   r:   rN   r9   )rP   detail_messager-   r-   r.   
_GetPrefix   s   
z_BaseProgressTracker._GetPrefixc                    D    fdd}zt  t j| _d _W dS  ty!   d _Y dS w )-Sets up a signal handler for handling SIGINT.c                    sJ    j r	t j j tjd W d    d S 1 sw   Y  d S Nz'

This operation cannot be cancelled.

)rH   r   OperationCancelledErrorrI   r@   r4   r5   writeunused_signalunused_framerU   r-   r.   _CtrlCHandler   s
   "z?_BaseProgressTracker._SetUpSignalHandler.<locals>._CtrlCHandlerTFNsignalSIGINTrJ   _restore_old_handler
ValueErrorrP   rd   r-   rU   r.   _SetUpSignalHandler   s   z(_BaseProgressTracker._SetUpSignalHandlerc                 C   6   | j rzttj| j W d S  ty   Y d S w d S rT   rh   rf   rg   rJ   ri   rU   r-   r-   r.   _TearDownSignalHandler      z+_BaseProgressTracker._TearDownSignalHandlerc                    sN       tj       jr% fdd}tj|d _	 j	
   S )Nc                         	 t  j   rd S qrT   
_SleepSecsrA   Tickr-   rU   r-   r.   Ticker   
   
z._BaseProgressTracker.__enter__.<locals>.Tickertarget)rk   r   file_only_loggerinfor[   _PrintrV   r>   ThreadrB   startrP   rt   r-   rU   r.   	__enter__   s   
z_BaseProgressTracker.__enter__c                 C   s   | j 2 d| _|rt|tjr| d n| d n| js.| jr)| |   n| d W d    n1 s8w   Y  | jrE| j	  | 
  d S )NTzaborted by ctrl-c.
zfailed.
zdone.
)r@   r=   
isinstancer   r_   rz   r7   r;   rB   joinrn   )rP   unused_ex_type	exc_valueunused_tracebackr-   r-   r.   __exit__   s   

z_BaseProgressTracker.__exit__c                 C      dS Give a visual indication to the user that some progress has been made.

    Output is sent to sys.stderr. Nothing is shown if output is not a TTY.

    Returns:
      Whether progress has completed.
    Nr-   rU   r-   r-   r.   rs         	z_BaseProgressTracker.Tickr2   c                 C   r   )z9Prints an update containing message to the output stream.Nr-   rP   r!   r-   r-   r.   rz         z_BaseProgressTracker._PrintNr2   )__name__
__module____qualname____doc__rR   propertyrV   r[   rk   rn   r~   r   abcabstractmethodrs   rz   r-   r-   r-   r.   r0   c   s    #
	

r0   c                       sB   e Zd ZdZ fddZdd Zdd Zdd	 ZdddZ  Z	S )r   r1   c                    s   |    tt|  S rT   )_SetupOutputsuperr   r~   rU   	__class__r-   r.   r~      s   z _NormalProgressTracker.__enter__c                    s2    fdd}t  j _ jj j|d _d S )Nc                     s,    j r   } | r jr| S d|  d S d S rX   )r:   rN   )rZ   rU   r-   r.   _FormattedCallback   s   z?_NormalProgressTracker._SetupOutput.<locals>._FormattedCallback)r#   )r	   SimpleSuffixConsoleOutputr6   _console_output
AddMessager9   _console_message)rP   r   r-   rU   r.   r      s
   	z#_NormalProgressTracker._SetupOutputc                 C   sX   | j  | js|  jd7  _| |   W d   n1 sw   Y  | j  | jS )r   r   N)r@   r=   r<   rz   
_GetSuffixr6   flushrU   r-   r-   r.   rs      s   
z_NormalProgressTracker.Tickc                 C   s8   | j r| jd }| j d|  S | jj| jt| jj  S )N   .)r,   r<   rM   
spin_markslen)rP   num_dotsr-   r-   r.   r     s   
z!_NormalProgressTracker._GetSuffixr2   c                 C   s.   | j s| jsdS | j| j| | j  dS )\  Reprints the prefix followed by an optional message.

    If there is a multiline message, we print the full message and every
    time the Prefix Message is the same, we only reprint the last line to
    account for a different 'message'. If there is a new message, we print
    on a new line.

    Args:
      message: str, suffix of message
    N)r7   rF   r   UpdateMessager   UpdateConsoler   r-   r-   r.   rz     s   z_NormalProgressTracker._Printr   )
r   r   r   r   r~   r   rs   r   rz   __classcell__r-   r-   r   r.   r      s    r   c                   @   s"   e Zd ZdZdd ZdddZdS )	r   r1   c                 C   sR   | j  | js| d W d   | jS W d   | jS 1 s!w   Y  | jS )r   r   N)r@   r=   rz   rU   r-   r-   r.   rs   $  s   

z#_NonInteractiveProgressTracker.Tickr2   c                 C   s0   | j s| jsdS |  }| j|p|d  dS )r   N
)r7   rF   r[   r6   r`   )rP   r!   display_messager-   r-   r.   rz   1  s
   z%_NonInteractiveProgressTracker._PrintNr   )r   r   r   r   rs   rz   r-   r-   r-   r.   r   !  s    r   c                   @   s0   e Zd ZdZdd Zdd Zdd Zdd	 Zd
S )r   z,A Progress tracker that doesn't do anything.c                 C   s   || _ || _d| _d S )NF)rH   rI   r=   )rP   r&   r'   r-   r-   r.   rR   K  s   
zNoOpProgressTracker.__init__c                         fdd}t  t j| _ S )Nc                        j r	t jd S rT   rH   r   r_   rI   ra   rU   r-   r.   rd   Q     z4NoOpProgressTracker.__enter__.<locals>._CtrlCHandlerrf   rg   rJ   rj   r-   rU   r.   r~   P     zNoOpProgressTracker.__enter__c                 C   rS   rT   r=   rU   r-   r-   r.   rs   W     zNoOpProgressTracker.Tickc                 C      d| _ ttj| j d S NTr=   rf   rg   rJ   rP   exc_typeexc_valexc_tbr-   r-   r.   r   Z     zNoOpProgressTracker.__exit__N)r   r   r   r   rR   r~   rs   r   r-   r-   r-   r.   r   H  s    r   c                       s,   e Zd ZdZ fddZ fddZ  ZS )r   zA Progress tracker that only prints deterministic start and end points.

  No UX about tracking should be exposed here. This is strictly for being able
  to tell that the tracker ran, not what it actually looks like.
  c                    s(   t t| || |pd| _tj| _d S )Nr2   )r   r   rR   r8   r4   r5   r6   )rP   r!   r&   r'   r   r-   r.   rR   f  s   
z_StubProgressTracker.__init__c                    s^   |sd}nt |tjrd}nd}t r%| jtjtjj	| j
|dd  tt| |||S )NSUCCESSINTERRUPTEDFAILURE)r!   statusr   )r   r   r_   r   rE   r6   r`   
JsonUXStubUXElementTypePROGRESS_TRACKERr8   r   r   r   )rP   r   r   r   r   r   r-   r.   r   k  s   
z_StubProgressTracker.__exit__)r   r   r   r   rR   r   r   r-   r-   r   r.   r   _  s    r   c                 C   s   t |  dS )z>Sleep int or float seconds. For mocking sleeps in this module.N)timesleep)secondsr-   r-   r.   rr   z  s   rr         @皙?      N@c                 C   sH   t jjj }|t jjjjjks|t jjjjjkrt	 S t
| ||||S )a  A context manager for visual feedback during long-running completions.

  A completion that exceeds the timeout is assumed to be refreshing the cache.
  At that point the progress tracker displays '?', forks the cache operation
  into the background, and exits.  This gives the background cache update a
  chance finish.  After background_ttl more seconds the update is forcibly
  exited (forced to call exit rather than killed by signal) to prevent not
  responding updates from proliferating in the background.

  Args:
    ofile: The stream to write to.
    timeout: float, The amount of time in second to show the tracker before
      backgrounding it.
    tick_delay: float, The time in second between ticks of the spinner.
    background_ttl: float, The number of seconds to allow the completion to
      run in the background before killing it.
    autotick: bool, True to tick the spinner automatically.

  Returns:
    The completion progress tracker.
  )r   r   r   r   r   r   r   r   r   _NoOpCompletionProgressTracker _NormalCompletionProgressTracker)ofiletimeoutr%   background_ttlr"   r)   r-   r-   r.   CompletionProgressTracker  s   
r   c                   @   sp   e Zd ZdZdZdd Zdd Z		dd	d
Zdd ZdddZ	dd Z
edd Zedd Zedd ZdS )r   a  A context manager for visual feedback during long-running completions.

  A completion that exceeds the timeout is assumed to be refreshing the cache.
  At that point the progress tracker displays '?', forks the cache operation
  into the background, and exits.  This gives the background cache update a
  chance finish.  After background_ttl more seconds the update is forcibly
  exited (forced to call exit rather than killed by signal) to prevent not
  responding updates from proliferating in the background.
  	   c                 C   s>   |p|   | _|| _|| _|| _|| _d| _t 	 | _
d S )Nr   )
_GetStream_ofile_timeoutrA   *_NormalCompletionProgressTracker__autotick_background_ttlr<   r   rK   rL   rM   )rP   r   r   r%   r   r"   r-   r-   r.   rR     s   z)_NormalCompletionProgressTracker.__init__c                 C   s2   | j rttj| j| _ttj| j| j| _| S rT   )	rV   rf   SIGALRM_Spin_old_handler	setitimerITIMER_REALrA   _old_itimerrU   r-   r-   r.   r~     s   z*_NormalCompletionProgressTracker.__enter__NTc                 C   sF   | j rtjtjg| jR   ttj| j |  s!| d d S d S NrY   )	rV   rf   r   r   r   r   r   	_TimedOut
_WriteMark)rP   unused_typeunused_valuer   r-   r-   r.   r     s   z)_NormalCompletionProgressTracker.__exit__c                 C   s
   | j dk S )z"True if the tracker has timed out.r   )r   rU   r-   r-   r.   r        
z*_NormalCompletionProgressTracker._TimedOutc                 C   s   |  j d7  _ | | jj| j t| jj   |  j| j8  _|  s%dS | d t	 r3t
d ttj| j ttj| j| j d| _dS )z4Rotates the spinner one tick and checks for timeout.r   N?)r<   r   rM   r   r   r   rA   r   osfork_exitrf   r   _ExitBackgroundr   r   r   r   )rP   
unused_sigrc   r-   r-   r.   r     s   



z&_NormalCompletionProgressTracker._Spinc                 C   s(   | j r| j |d  | j   dS dS )zWrites one mark to self._ofile.N)r   r`   r   )rP   markr-   r-   r.   r     s   z+_NormalCompletionProgressTracker._WriteMarkc                   C   s   t d dS )zEUnconditionally exits the background completer process after timeout.r   N)r   r   r-   r-   r-   r.   r     s   z0_NormalCompletionProgressTracker._ExitBackgroundc                 C   rS   rT   )r   rU   r-   r-   r.   rV     rW   z*_NormalCompletionProgressTracker._autotickc                   C   s   t t tjdS )z$Returns the completer output stream.w)r   fdopendupr   _COMPLETION_FDr-   r-   r-   r.   r     s   z+_NormalCompletionProgressTracker._GetStream)NTNNN)r   r   r   r   r   rR   r~   r   r   r   r   staticmethodr   r   rV   r   r-   r-   r-   r.   r     s"    
	



r   c                   @   s(   e Zd ZdZdd Zdd Zdd ZdS )	r   z'A Progress tracker that prints nothing.c                 C      d S rT   r-   rU   r-   r-   r.   rR        z'_NoOpCompletionProgressTracker.__init__c                 C   s   | S rT   r-   rU   r-   r-   r.   r~     r   z(_NoOpCompletionProgressTracker.__enter__c                 C   r   rT   r-   r   r-   r-   r.   r     r   z'_NoOpCompletionProgressTracker.__exit__N)r   r   r   r   rR   r~   r   r-   r-   r-   r.   r     s
    r   c                 C   s   t jjj }|s|t jjjjjkrt|||
S |t jjjj	jkr(t
| |||
S tjdd}|r<t  r9t}nt}nt}|| ||||	||||
||S )a  A progress tracker for performing actions with multiple stages.

  The progress tracker is a context manager. To start displaying information
  about a running stage, call StartStage within the staged progress tracker
  context. To update the message of a stage, use UpdateStage. When a stage is
  completed/failed there are CompleteStage and FailStage methods on the
  tracker as well.

  Note that stages do not need to be started/completed in order. In
  non-multiline (the only supported mode) output mode, the displayed stage will
  be the earliest started stage that has not been completed.

  Example Usage:
    stages = [
      Stage('Getting bread...', key='bread'),
      Stage('Getting peanut butter...', key='pb'),
      Stage('Making sandwich...', key='make')]
    with StagedProgressTracker(
        'Making sandwich...',
        stages,
        success_message='Time to eat!',
        failure_message='Time to order delivery..!',
        tracker_id='meta.make_sandwich') as tracker:
      tracker.StartStage('bread')
      # Go to pantry
      tracker.UpdateStage('bread', 'Looking for bread in the pantry')
      # Get bread
      tracker.CompleteStage('bread', 'Got some whole wheat bread!')

      tracker.StartStage('pb')
      # Look for peanut butter
      if pb_not_found:
        error = exceptions.NoPeanutButterError('So sad!')
        tracker.FailStage('pb', error)
      elif pb_not_organic:
        tracker.CompleteStageWithWarning('pb', 'The pb is not organic!')
      else:
        tracker.CompleteStage('bread', 'Got some organic pb!')

  Args:
    message: str, The message to show next to the spinner.
    stages: list[Stage], A list of stages for the progress tracker to run. Once
      you pass the stages to a StagedProgressTracker, they're owned by the
      tracker and you should not mutate them.
    tracker_id: str The ID of this tracker that will be used for metrics.
    autotick: bool, True to have the spinner tick on its own. Otherwise, you
      need to call Tick() explicitly to move the spinner.
    tick_delay: float, The amount of time to wait between ticks, in second.
    interruptable: boolean, True if the user can ctrl-c the operation. If so,
      it will stop and will report as aborted. If False,
    done_message_callback: func, A callback to get a more detailed done message.
    success_message: str, A message to display on success of all tasks.
    warning_message: str, A message to display when no task fails but one or
      more tasks complete with a warning and none fail.
    failure_message: str, A message to display on failure of a task.
    aborted_message: str, A custom message to put in the exception when it is
      cancelled by the user.
    suppress_output: bool, True to suppress output from the tracker.

  Returns:
    The progress tracker.
  Tr   )r   r   r   r   r   r   r   r   NoOpStagedProgressTrackerr   _StubStagedProgressTrackerr   r   r   rC   SupportsAnsi_MultilineStagedProgressTracker_NormalStagedProgressTracker$_NonInteractiveStagedProgressTracker)r!   stages
tracker_idr"   r%   r&   r$   success_messagewarning_messagefailure_messager'   suppress_outputr)   r*   r+   r-   r-   r.   StagedProgressTracker  s&   D

r   c                   @   s>   e Zd ZdZdddZedd Zedd Zed	d
 ZdS )Stagez-Defines a stage of a staged progress tracker.Nc                 C   s8   || _ |dur	|n| j | _d| _|| _d| _tj| _dS )aZ  Encapsulates a stage in a staged progress tracker.

    A task should contain a message about what it does.

    Args:
      header: (str) The header that describes what the task is doing.
        A high level description like 'Uploading files' would be appropriate.
      key: (str) A key which can be used to access/refer to this stage. Must be
        unique within a StagedProgressTracker. If not provided, the header will
        be used as the key.
      task_id: (str) The ID of this task that will be used for metrics.
      timing metrics. NOTE: Metrics are currently not implemented yet.
    Nr2   F)_header_keyr!   task_id_is_doneStageCompletionStatusNOT_STARTEDr   )rP   headerkeyr   r-   r-   r.   rR   _  s   zStage.__init__c                 C   rS   rT   )r   rU   r-   r-   r.   r  x  rW   z	Stage.keyc                 C   rS   rT   )r   rU   r-   r-   r.   r  |  rW   zStage.headerc                 C   rS   rT   )r   rU   r-   r-   r.   is_done  rW   zStage.is_doner   )	r   r   r   r   rR   r   r  r  r  r-   r-   r-   r.   r   \  s    


r   c                   @   s(   e Zd ZdZdZdZdZdZdZdZ	dS )	r   z+Indicates the completion status of a stage.znot startedzstill runningdonefailedinterruptedwarningN)
r   r   r   r   r  RUNNINGr   FAILEDr   WARNINGr-   r-   r-   r.   r     s    r   c                   @   s4  e Zd ZdZ	dFddZdd Zdd Zd	d
 Zedd Z	dd Z
dd Zdd Zdd Zdd Zdd Zdd Zdd Zdd Zdd  Zejd!d" Zd#d$ Zejd%d& Zd'd( Zd)d* ZdGd,d-Zd.d/ Zd0d1 Zd2d3 ZdHd4d5Zd6d7 Z dFd8d9Z!d:d; Z"d<d= Z#d>d? Z$d@dA Z%dFdBdCZ&dDdE Z'dS )I_BaseStagedProgressTrackera{  Base class for staged progress trackers.

  During each tick, the tracker checks if there is a stage being displayed by
  checking if _stage_being_displayed is not None. If it is not none and stage
  has not completed, then the tracker will print an update. If the stage is
  done, then the tracker will write out the status of all completed stages
  in _running_stages_queue.
  Nc                 C   s  t  | _|D ]}|j| jv rtd|j|| j|j< qtj| _|| _	|| _
|| _|| _|	| _|| _|
| _|d u r?t }| d }t|trN|dk rPd}t oW|dk| _|o]| j| _|| _|| _| | _d| _d| _d| _d | _ t! | _"g | _#g | _$g | _%t&' | _(d S )NzDuplicate stage key: {}r   FT))collectionsOrderedDict_stagesr  ri   formatr4   r5   r6   r8   _success_message_warning_message_failure_messagerI   r;   _tracker_idr   rK   rD   r   intr   rE   rF   $_BaseStagedProgressTracker__autotickrH   rA   rL   rM   r=   _exception_is_uncaughtr<   rB   set_running_stages_completed_stages_completed_with_warnings_stages_exit_output_warningsr>   r?   r@   )rP   r!   r   r   r   r   r"   r%   r&   r'   r   r$   consolestagerQ   r-   r-   r.   rR     s@   

z#_BaseStagedProgressTracker.__init__c                 C   s
   | j | S rT   )r  )rP   r  r-   r-   r.   __getitem__     
z&_BaseStagedProgressTracker.__getitem__c                 C   
   t | jS rT   )iterr  rU   r-   r-   r.   __iter__  r   z#_BaseStagedProgressTracker.__iter__c                 C   r!  rT   )r   r  rU   r-   r-   r.   __len__  r   z"_BaseStagedProgressTracker.__len__c                 C   rS   rT   )r  rU   r-   r-   r.   rV     rW   z$_BaseStagedProgressTracker._autotickc                 C   s   |  |p	| | S )z&Returns True if the stage is complete.)	IsRunning	IsWaitingrP   r  r-   r-   r.   
IsComplete  s   z%_BaseStagedProgressTracker.IsCompletec                 C      | j |dd}|jtjkS )z%Returns True if the stage is running.Tallow_complete)_ValidateStager   r   r	  r'  r-   r-   r.   r%       z$_BaseStagedProgressTracker.IsRunningc                 C   r!  )zBReturns True if this tracker has encountered at least one warning.)boolr  rU   r-   r-   r.   
HasWarning  r   z%_BaseStagedProgressTracker.HasWarningc                 C   r)  )z-Returns True if the stage is not yet started.Tr*  )r,  r   r   r  r'  r-   r-   r.   r&    r-  z$_BaseStagedProgressTracker.IsWaitingc                    r\   )r]   c                    s    j r	t j   d S rT   )rH   r   r_   rI   _NotifyUninterruptableErrorra   rU   r-   r.   rd     s   zE_BaseStagedProgressTracker._SetUpSignalHandler.<locals>._CtrlCHandlerTFNre   rj   r-   rU   r.   rk     s   z._BaseStagedProgressTracker._SetUpSignalHandlerc                 C   s8   | j  tjd W d    d S 1 sw   Y  d S r^   )r@   r4   r5   r`   rU   r-   r-   r.   r0    s   "z6_BaseStagedProgressTracker._NotifyUninterruptableErrorc                 C   rl   rT   rm   rU   r-   r-   r.   rn     ro   z1_BaseStagedProgressTracker._TearDownSignalHandlerc                    s\          tj j     jr, fdd}tj	|d _
d j
_ j
   S )Nc                      rp   rT   rq   r-   rU   r-   r.   rt     ru   z4_BaseStagedProgressTracker.__enter__.<locals>.Tickerrv   T)r   rk   r   rx   ry   r8   rz   rV   r>   r{   rB   daemonr|   r}   r-   rU   r.   r~     s   
z$_BaseStagedProgressTracker.__enter__c                 C   s   | j  d| _|r| jr| | n| j|  d W d    n1 s$w   Y  | jr1| j  |   | j	D ]
}t
jd|  q8d S )NT)warnedz  %s)r@   r=   r  _HandleUncaughtException_PrintExitOutputr/  rB   r   rn   r  r   r   Print)rP   r   r   r   r   r-   r-   r.   r     s   
	

z#_BaseStagedProgressTracker.__exit__c                 C   s,   t |tjr| jdd d S | jdd d S )NTabortedr  )r   r   r_   r4  rP   r   r-   r-   r.   r3    s   z3_BaseStagedProgressTracker._HandleUncaughtExceptionc                 C   r   )zASets up the output for the tracker. Gets called during __enter__.Nr-   rU   r-   r-   r.   r   '  r   z'_BaseStagedProgressTracker._SetupOutputc                 C   r   )z)Updates the header messsage if supported.Nr-   r   r-   r-   r.   UpdateHeaderMessage,     z._BaseStagedProgressTracker.UpdateHeaderMessagec                 C   r   r   r-   rU   r-   r-   r.   rs   0  r   z_BaseStagedProgressTracker.Tickc                 C   s   | j j| jt| j j  S )Returns the next tick mark.)rM   r   r<   r   rP   ticksr-   r-   r.   _GetTickMark;  s   z'_BaseStagedProgressTracker._GetTickMarkc                 C   s   |j S rT   valuerP   r   r-   r-   r.   _GetStagedCompletedSuffix?  r   z4_BaseStagedProgressTracker._GetStagedCompletedSuffixFc                 C   s<   || vrt d| |}|s|jtjtjhvrt d|S )zValidates the stage belongs to the tracker.

    Args:
      key: the key of the stage to validate.
      allow_complete: whether to error on an already-complete stage

    Returns:
      The validated Stage object, even if we were passed a key.
    z4This stage does not belong to this progress tracker.z!This stage has already completed.)ri   getr   r   r  r	  )rP   r  r+  r  r-   r-   r.   r,  B  s   

z)_BaseStagedProgressTracker._ValidateStagec                 C   sZ   |  |}| j | j| tj|_| | W d   n1 s"w   Y  |   dS z9Informs the progress tracker that this stage has started.N)	r,  r@   r  addr   r	  r   _StartStagers   )rP   r  r  r-   r-   r.   
StartStageT  s   
z%_BaseStagedProgressTracker.StartStagec                 C   r   )z3Override to customize behavior on starting a stage.Nr-   r'  r-   r-   r.   rG  ]  r;  z&_BaseStagedProgressTracker._StartStagec                 C   r   )z2Override to customize behavior on failing a stage.Nr-   )rP   r  failure_exceptionr!   r-   r-   r.   
_FailStagea  r;  z%_BaseStagedProgressTracker._FailStagec                 C   r   )z7Override to customize behavior on printing exit output.Nr-   )rP   r7  r2  r  r-   r-   r.   r4  e  r;  z+_BaseStagedProgressTracker._PrintExitOutputc                 C   sB   |  |}| j ||_W d   n1 sw   Y  |   dS )z(Updates a stage in the progress tracker.N)r,  r@   r!   rs   rP   r  r!   r  r-   r-   r.   UpdateStagei  s
   
z&_BaseStagedProgressTracker.UpdateStagec                 C   sn   |  |}| j! tj|_d|_| j| |dur||_| 	| W d   n1 s,w   Y  | 
  dS )z;Informs the progress tracker that this stage has completed.TN)r,  r@   r   r   r   r   r  discardr!   _CompleteStagers   rK  r-   r-   r.   CompleteStageq  s   
z(_BaseStagedProgressTracker.CompleteStagec                 C   r   rT   r-   r'  r-   r-   r.   rN  }  r   z)_BaseStagedProgressTracker._CompleteStagec                 C   s   |  ||g d S rT   )CompleteStageWithWarnings)rP   r  r   r-   r-   r.   CompleteStageWithWarning     z3_BaseStagedProgressTracker.CompleteStageWithWarningc                 C   s|   |  |}| j( tj|_d|_| j| | j	| | j
|j | || W d   n1 s3w   Y  |   dS )zInforms the progress tracker that this stage completed with warnings.

    Args:
      key: str, key for the stage to fail.
      warning_messages: list of str, user visible warning messages.
    TN)r,  r@   r   r  r   r   r  rM  r  extendr  appendr  _CompleteStageWithWarningsrs   )rP   r  warning_messagesr  r-   r-   r.   rP    s   
z4_BaseStagedProgressTracker.CompleteStageWithWarningsc                 C   r   )zCOverride to customize behavior on completing a stage with warnings.Nr-   rP   r  rV  r-   r-   r.   rU    r;  z5_BaseStagedProgressTracker._CompleteStageWithWarningsc                 C   s   |  |}| j# tj|_d|_| j| |dur||_| 	||| W d   n1 s.w   Y  | 
  |rD| jdd d| _|dS )zInforms the progress tracker that this stage has failed.

    Args:
      key: str, key for the stage to fail.
      failure_exception: Exception, raised by __exit__.
      message: str, user visible message for failure.
    TNr8  F)r,  r@   r   r
  r   r   r  rM  r!   rJ  rs   r4  r  )rP   r  rI  r!   r  r-   r-   r.   	FailStage  s   
z$_BaseStagedProgressTracker.FailStagec                 C   s8   | j  | j| W d   dS 1 sw   Y  dS )zAdd a warning message independent of any particular stage.

    This warning message will be printed on __exit__.

    Args:
      warning_message: str, user visible warning message.
    N)r@   r  rT  )rP   r   r-   r-   r.   
AddWarning  s   "z%_BaseStagedProgressTracker.AddWarningrT   )FFFF)(r   r   r   r   rR   r  r#  r$  r   rV   r(  r%  r/  r&  rk   r0  rn   r~   r   r3  r   r   r   r:  rs   r?  rC  r,  rH  rG  rJ  r4  rL  rO  rN  rQ  rP  rU  rX  rY  r-   r-   r-   r.   r    sN    
)





	


r  c                       sv   e Zd ZdZ fddZdd ZdddZd	d
 Zdd Zdd Z	dddZ
dd Zdd Zdd ZdddZ  ZS )r     A context manager for telling the user about long-running progress.

  This class uses the core.console.multiline.ConsoleOutput interface for
  outputting. The header and each stage is defined as a message object
  contained by the ConsoleOutput message.
  c                    s&   g | _ d | _tt| j|i | d S rT   )_running_stages_queue_stage_being_displayedr   r   rR   rP   argskwargsr   r-   r.   rR     s   z%_NormalStagedProgressTracker.__init__c                 C   s*   t | j| _| j| j| _| j| _d S rT   )r	   r   r6   r   r   r8   _header_message_current_stage_messagerU   r-   r-   r.   r     s   z)_NormalStagedProgressTracker._SetupOutputNc                 C   s&   | j D ]}||krtj|_d|_qd S r   )r\  r   r   r   r   )rP   r  rI  r!   running_stager-   r-   r.   rJ    s
   
z'_NormalStagedProgressTracker._FailStagec                 C   s&   | j | | jdu r|   dS dS rE  )r\  rT  r]  _LoadNextStageForDisplayr'  r-   r-   r.   rG    s   
z(_NormalStagedProgressTracker._StartStagec                 C   s&   | j r| j d | _| | j dS d S )Nr   T)r\  r]  _SetUpOutputForStagerU   r-   r-   r.   rd    s
   z5_NormalStagedProgressTracker._LoadNextStageForDisplayc                 C   s   | j n | js_|  jd7  _| jdu r|   n2| jrJ| jd jrJ| jd}| j	|j
 | | jj}| | |  sAd| _| jrJ| jd js!| jrh| | | j W d   | jS W d   | jS W d   | jS 1 stw   Y  | jS )"  Give a visual indication to the user that some progress has been made.

    Output is sent to sys.stderr. Nothing is shown if output is not a TTY.
    This method also handles loading new stages and flushing out completed
    stages.

    Returns:
      Whether progress has completed.
    r   Nr   )r@   r=   r<   r]  rd  r\  r  popr  rT  r  rC  r   rz   r?  )rP   completed_stagecompletion_statusr-   r-   r.   rs     s>   









z!_NormalStagedProgressTracker.TickFc                 C   sj   |    |r| jp
d}n|r| jpd}n|r| jpd}n| jp d}| jr,|d|   7 }| |d  dS )2Handles the final output for the progress tracker.Aborted.Failed.Completed with warnings:Done.rY   r   N)_SetupExitOutputrI   r  r  r  r;   rz   )rP   r7  r2  r  msgr-   r-   r.   r4    s   
z-_NormalStagedProgressTracker._PrintExitOutputc                 C   s   | j d| _dS )-Sets up output to print out the closing line.r2   N)r   r   rb  rU   r-   r-   r.   ro    s   z-_NormalStagedProgressTracker._SetupExitOutputc                 C   sH   t |tjr| d | jdd d S | | tj | jdd d S )Nzaborted by ctrl-cTr6  r8  )r   r   r_   rz   r4  rC  r   r
  r9  r-   r-   r.   r3    s   

z5_NormalStagedProgressTracker._HandleUncaughtExceptionc                    s&    fdd}| j j jd|d| _d S )Nc                      s    j r
d j  d S d S rX   )r!   r-   r  r-   r.   r      s   zM_NormalStagedProgressTracker._SetUpOutputForStage.<locals>._FormattedCallbackr   )indentation_levelr#   )r   r   r  rb  )rP   r  r   r-   rr  r.   re    s   z1_NormalStagedProgressTracker._SetUpOutputForStager2   c                 C   s2   | j sdS | jr| j| j| | j  dS dS zoPrints an update containing message to the output stream.

    Args:
      message: str, suffix of message
    N)rF   rb  r   r   r   r   r-   r-   r.   rz   )  s   z#_NormalStagedProgressTracker._PrintrT   rZ  r   )r   r   r   r   rR   r   rJ  rG  rd  rs   r4  ro  r3  re  rz   r   r-   r-   r   r.   r     s    

 
r   c                   @   sB   e Zd ZdZdd Zdd Zdd Zdd	 Zd
d ZdddZ	dS )r   r1   c                 C   r   )rq  Nr-   rU   r-   r-   r.   ro  9  r   z5_NonInteractiveStagedProgressTracker._SetupExitOutputc                 C   s   |  | jd  d S Nr   )rz   r8   rU   r-   r-   r.   r   >  s   z1_NonInteractiveStagedProgressTracker._SetupOutputc                 C   r   )r<  r   r-   r=  r-   r-   r.   r?  A  r;  z1_NonInteractiveStagedProgressTracker._GetTickMarkc                 C   s
   |j d S ru  r@  rB  r-   r-   r.   rC  E  r   z>_NonInteractiveStagedProgressTracker._GetStagedCompletedSuffixc                 C   s,   |j }|jr|d|j d 7 }| | d S rX   )r  r!   rz   rP   r  r!   r-   r-   r.   re  H  s   z9_NonInteractiveStagedProgressTracker._SetUpOutputForStager2   c                 C   s   | j sdS | j| dS rt  )rF   r6   r`   r   r-   r-   r.   rz   N  s   z+_NonInteractiveStagedProgressTracker._PrintNr   )
r   r   r   r   ro  r   r?  rC  re  rz   r-   r-   r-   r.   r   6  s    r   c                       s   e Zd ZdZ fddZdd Zdd Zd&d	d
Zdd Zd'ddZ	dd Z
dd Zd(ddZdd Zdd Zd)ddZdd Zd d! Zd"d# Zd&d$d%Z  ZS )*r   r[  c                    s$   t  | _tt| j|i | d S rT   )r
   GetTypedTextParser_parserr   r   rR   r^  r   r-   r.   rR   a  s   
z(_MultilineStagedProgressTracker.__init__c                 C   s   || j _d S rT   )_header_stager!   r   r-   r-   r.   r:  e  s   z3_MultilineStagedProgressTracker.UpdateHeaderMessagec                 C   s4   || j  }| jjr|d| jj 7 }| | j| d S r   )r8   ry  r!   _UpdateMessagera  )rP   prefixr!   r-   r-   r.   _UpdateHeaderMessagek  s   
z4_MultilineStagedProgressTracker._UpdateHeaderMessager2   c                 C   sB   |  |j|}|j}|jr|d|j 7 }| | j| ||  d S r   )_GenerateStagePrefixr   r  r!   rz  _stage_messages)rP   r  	tick_markr{  r!   r-   r-   r.   _UpdateStageTickMarkq  s
   z4_MultilineStagedProgressTracker._UpdateStageTickMarkc                 C   s   | j |}| j|| d S rT   )rx  ParseTypedTextToStringr   r   rv  r-   r-   r.   rz  x  s   z._MultilineStagedProgressTracker._UpdateMessager   c                 C   s   | j |}| jj||dS )Nrs  )rx  r  r   r   )rP   r!   rs  r-   r-   r.   _AddMessage|  s   z+_MultilineStagedProgressTracker._AddMessagec                 C   s<   | j  | d W d    n1 sw   Y  |   d S )Nz#This operation cannot be cancelled.)r@   r:  rs   rU   r-   r-   r.   r0    s   z;_MultilineStagedProgressTracker._NotifyUninterruptableErrorc                 C   s   | j dS )rq  r2   )r   r   rU   r-   r-   r.   ro    s   z0_MultilineStagedProgressTracker._SetupExitOutputFc                 C   s   |   }|r| jp
d}tj| j_n$|r| jpd}tj| j_n|r+| jp$d}tj| j_n
| jp/d}tj	| j_| j
r@|d| 
  7 }| || | | jj dS )rj  rk  rl  rm  rn  rY   N)ro  rI   r   r
  ry  r   r  r  r  r   r;   rz  rz   rM   r  )rP   r7  r2  r  output_messagerp  r-   r-   r.   r4    s    




z0_MultilineStagedProgressTracker._PrintExitOutputc                 C   s|   d| _ t| j| _| | j| _td| _	t
j| j	_t | _| j D ]}| j|jdd| j|< | | q$| j  d S )NFr2   r   r  )_maintain_queuer	   MultilineConsoleOutputr6   r   r  r8   ra  r   ry  r   r	  r   dictr~  r  valuesr  r  r   r'  r-   r-   r.   r     s   

z,_MultilineStagedProgressTracker._SetupOutputc                 C   sf   |t jkr
| jj}n|t jkr| jj}n|t jkr| jj}n	|t jkr'| jj	}|d| jj
t|   S r   )r   r  rM   not_startedr   successr
  r  r   r  prefix_lengthr   )rP   stage_statusr  r-   r-   r.   r}    s   






z4_MultilineStagedProgressTracker._GenerateStagePrefixNc                 C   sH   |  | |r | j D ]}||kr|jtjkrtj|_d|_qdS dS )z8Informs the progress tracker that this stage has failed.TN)r  r  r  r   r   r	  r   r   )rP   r  	exceptionr!   other_stager-   r-   r.   rJ    s   
z*_MultilineStagedProgressTracker._FailStagec                 C      |  | d S rT   r  r'  r-   r-   r.   rN       z._MultilineStagedProgressTracker._CompleteStagec                 C   r  rT   r  rW  r-   r-   r.   rU    r  z:_MultilineStagedProgressTracker._CompleteStageWithWarningsc                 C   sh   | j & | js |  jd7  _| | | j W d   | jS W d   | jS 1 s,w   Y  | jS )rf  r   N)r@   r=   r<   rz   r?  rU   r-   r-   r.   rs     s   


z$_MultilineStagedProgressTracker.Tickc                 C   sN   | j sdS | | jj|}| | | jD ]
}| | | | q| j  dS )zqPrints an update containing message to the output stream.

    Args:
      tick_mark: str, suffix of message
    N)	rF   r}  ry  r   r|  r  r  r   r   )rP   r  header_prefixr  r-   r-   r.   rz     s   

z&_MultilineStagedProgressTracker._Printr   )r   rZ  rT   )r   r   r   r   rR   r:  r|  r  rz  r  r0  ro  r4  r   r}  rJ  rN  rU  rs   rz   r   r-   r-   r   r.   r   Y  s$    




r   c                       sT   e Zd ZdZd fdd	Zdd Zddd	Zd
d Zdd Zdd Z	dd Z
  ZS )r   z3A staged progress tracker that doesn't do anything.Fr2   c                    sB   t t| jd|ddddd||dd tjdddd || _d| _d S )Nr2   Fr   asciiT)encodingr   )r!   r   r   r   r   r"   r%   r&   r'   r   r$   r  )r   r   rR   r   rC   rI   r=   )rP   r   r&   r'   r   r-   r.   rR     s$   

z"NoOpStagedProgressTracker.__init__c                    r   )Nc                    r   rT   r   ra   rU   r-   r.   rd     r   z:NoOpStagedProgressTracker.__enter__.<locals>._CtrlCHandlerr   rj   r-   rU   r.   r~     r   z#NoOpStagedProgressTracker.__enter__c                 C   r   rT   r-   r   r-   r-   r.   rz     r;  z NoOpStagedProgressTracker._Printc                 C   rS   rT   r   rU   r-   r-   r.   rs   	  r   zNoOpStagedProgressTracker.Tickc                 C   r   r   r   r   r-   r-   r.   r     r   z"NoOpStagedProgressTracker.__exit__c                 C   r   rT   r-   rU   r-   r-   r.   r     r   z&NoOpStagedProgressTracker._SetupOutputc                 C   r   rT   r-   r   r-   r-   r.   r:    r   z-NoOpStagedProgressTracker.UpdateHeaderMessage)Fr2   r   )r   r   r   r   rR   r~   rz   rs   r   r   r:  r   r-   r-   r   r.   r     s    
r   c                       s>   e Zd ZdZ fddZdd ZdddZ fd	d
Z  ZS )r   zStaged tracker that only prints deterministic start and end points.

  No UX about tracking should be exposed here. This is strictly for being able
  to tell that the tracker ran, not what it actually looks like.
  c                    s2   t t| ||| || _g | _d | _tj| _d S rT   )	r   r   rR   r8   _succeeded_stages_failed_stager4   r5   r6   )rP   r!   r   r&   r'   r   r-   r.   rR     s   
z#_StubStagedProgressTracker.__init__c                 C   s   | j |j d S rT   )r  rT  r  r'  r-   r-   r.   rN  &  rR  z)_StubStagedProgressTracker._CompleteStageNc                 C   s   |j | _|rT   )r  r  )rP   r  r  r!   r-   r-   r.   rJ  )  s   z%_StubStagedProgressTracker._FailStagec              	      sx   |rt |tjrd}n|rd}n	|  rd}nd}t r2| jtjtj	j
| j|| j| jdd  tt| |||S )Nr   r   r  r   )r!   r   succeeded_stagesfailed_stager   )r   r   r_   r/  r   rE   r6   r`   r   r   STAGED_PROGRESS_TRACKERr8   r  r  r   r   r   )rP   r   r   r   status_messager   r-   r.   r   -  s,   

z#_StubStagedProgressTracker.__exit__rT   )	r   r   r   r   rR   rN  rJ  r   r   r-   r-   r   r.   r     s    
r   )Nr   r   r   T)3r   
__future__r   r   r   r   r  r   rf   r4   r>   r   enumgooglecloudsdk.corer   r   googlecloudsdk.core.consoler   r   r	   !googlecloudsdk.core.console.styler
   sixcollections_abcversion_infor_   DEFAULT_MESSAGEr/   with_metaclassABCMetaobjectr0   r   r   r   r   rr   r   r   r   r   r   Enumr   Mappingr  r   r   r   r   r   r-   r-   r-   r.   <module>   sz   

6 >'
!U
Z)
  -|# .