1616import shlex
1717import subprocess
1818from argparse import (
19+ Action ,
1920 ArgumentDefaultsHelpFormatter ,
2021 ArgumentParser ,
2122 Namespace ,
@@ -74,6 +75,23 @@ class StagedScript:
7475 retry_arg_group (argparse._ArgumentGroup): A container within
7576 the :class:`ArgumentParser` holding all the arguments
7677 associated with retrying stages.
78+ retry_attempts (Dict[str, int]): A mapping from stage names to
79+ the number of times to attempt retrying each stage.
80+ retry_attempts_arg (Dict[str, argparse.Action]): The
81+ corresponding arguments in the :class:`ArgumentParser`, so
82+ subclass developers can modify them if needed.
83+ retry_delay (Dict[str, float]): A mapping from stage names to
84+ how long to wait (in seconds) before attempting to retry a
85+ stage.
86+ retry_delay_arg (Dict[str, argparse.Action]): The
87+ corresponding arguments in the :class:`ArgumentParser`, so
88+ subclass developers can modify them if needed.
89+ retry_timeout (Dict[str, int]): A mapping from stage names to
90+ how long to wait (in seconds) before giving up on retrying
91+ the stage.
92+ retry_timeout_arg (Dict[str, argparse.Action]): The
93+ corresponding arguments in the :class:`ArgumentParser`, so
94+ subclass developers can modify them if needed.
7795 script_name (str): The name of the script (the
7896 :class:`StagedScript` subclass) being run.
7997 script_stem (str): Same as :attr:`script_name`, but without the
@@ -90,28 +108,6 @@ class StagedScript:
90108 the user via the command line arguments.
91109 start_time (datetime): The time at which this object was
92110 initialized.
93-
94- Note that additional attributes are automatically generated for each
95- ``stage`` registered for a subclass object:
96-
97- ``STAGE_NAME_retry_attempts`` (int)
98- The number of times to attempt retrying the ``STAGE_NAME``
99- stage.
100- ``STAGE_NAME_retry_attempts_arg`` (argparse.Action)
101- The corresponding argument in the :class:`ArgumentParser`, so
102- subclass developers can modify it if needed.
103- ``STAGE_NAME_retry_delay`` (float)
104- How long to wait (in seconds) before attempting to retry the
105- ``STAGE_NAME`` stage.
106- ``STAGE_NAME_retry_delay_arg`` (argparse.Action)
107- The corresponding argument in the :class:`ArgumentParser`, so
108- subclass developers can modify it if needed.
109- ``STAGE_NAME_retry_timeout`` (int)
110- How long to wait (in seconds) before giving up on retrying the
111- ``STAGE_NAME`` stage.
112- ``STAGE_NAME_retry_timeout_arg`` (argparse.Action)
113- The corresponding argument in the :class:`ArgumentParser`, so
114- subclass developers can modify it if needed.
115111 """
116112
117113 def __init__ (
@@ -155,6 +151,12 @@ def __init__(
155151 self .dry_run = False
156152 self .durations : List [StageDuration ] = []
157153 self .print_commands = print_commands
154+ self .retry_attempts : Dict [str , int ] = {}
155+ self .retry_attempts_arg : Dict [str , Action ] = {}
156+ self .retry_delay : Dict [str , float ] = {}
157+ self .retry_delay_arg : Dict [str , Action ] = {}
158+ self .retry_timeout : Dict [str , int ] = {}
159+ self .retry_timeout_arg : Dict [str , Action ] = {}
158160 self .script_name = Path (__main__ .__file__ ).name
159161 self .script_stem = Path (__main__ .__file__ ).stem
160162 self .script_success = True
@@ -320,9 +322,9 @@ def wrapper(self, *args, **kwargs) -> None:
320322 """
321323 self .current_stage = stage_name
322324 get_phase_method (self , "_run_pre_stage_actions" )()
323- timeout = getattr ( self , f" { stage_name } _retry_timeout" )
324- attempts = getattr ( self , f" { stage_name } _retry_attempts" )
325- delay = getattr ( self , f" { stage_name } _retry_delay" )
325+ timeout = self . retry_timeout [ stage_name ]
326+ attempts = self . retry_attempts [ stage_name ]
327+ delay = self . retry_delay [ stage_name ]
326328 stop_after_timeout = stop_after_delay (timeout )
327329 stop_after_max_attempts = stop_after_attempt (attempts + 1 )
328330 retry = Retrying (
@@ -590,10 +592,7 @@ def _handle_stage_retry_error_test(
590592 retry: The :class:`Retrying` controller, which contains
591593 information about the retrying that was done.
592594 """
593- retry_attempts = getattr (
594- self , f"{ self .current_stage } , 0
595- )
596- if retry_attempts > 0 :
595+ if self .retry_attempts [self .current_stage ] > 0 :
597596 stage_time = timedelta (
598597 seconds = retry .statistics ["delay_since_first_attempt" ]
599598 )
@@ -660,9 +659,9 @@ def parser(self) -> ArgumentParser:
660659
661660 .. code-block:: python
662661
663- self.foo_retry_attempts_arg .help = argparse.SUPPRESS
664- self.foo_retry_delay_arg .help = argparse.SUPPRESS
665- self.foo_retry_timeout_arg .help = argparse.SUPPRESS
662+ self.retry_attempts_arg["foo"] .help = argparse.SUPPRESS
663+ self.retry_delay_arg["foo"] .help = argparse.SUPPRESS
664+ self.retry_timeout_arg["foo"] .help = argparse.SUPPRESS
666665
667666 And if you want to remove the title for the retry group
668667 altogether, you can do so with:
@@ -706,23 +705,23 @@ def parser(self) -> ArgumentParser:
706705 type = int ,
707706 help = f"How many times to retry the { stage !r} ,
708707 )
709- setattr ( self , f" { stage } _retry_attempts_arg" , retry_attempts )
708+ self . retry_attempts_arg [ stage ] = retry_attempts
710709 retry_delay = self .retry_arg_group .add_argument (
711710 f"--{ stage } ,
712711 default = 0 ,
713712 type = float ,
714713 help = "How long to wait (in seconds) before retrying the "
715714 f"{ stage !r} ,
716715 )
717- setattr ( self , f" { stage } _retry_delay_arg" , retry_delay )
716+ self . retry_delay_arg [ stage ] = retry_delay
718717 retry_timeout = self .retry_arg_group .add_argument (
719718 f"--{ stage } ,
720719 default = 60 ,
721720 type = int ,
722721 help = "How long to wait (in seconds) before giving up on "
723722 f"retrying the { stage !r} ,
724723 )
725- setattr ( self , f" { stage } _retry_timeout_arg" , retry_timeout )
724+ self . retry_timeout_arg [ stage ] = retry_timeout
726725 return my_parser
727726
728727 def parse_args (self , argv : List [str ]) -> None :
@@ -755,12 +754,13 @@ def parse_args(self, argv: List[str]) -> None:
755754 set (self .args .stage ) if self .args .stage is not None else set ()
756755 )
757756 for stage in self .stages :
758- for retry_arg in [
759- f"{ stage } ,
760- f"{ stage } ,
761- f"{ stage } ,
762- ]:
763- setattr (self , retry_arg , getattr (self .args , retry_arg , None ))
757+ for arg in ["attempts" , "delay" , "timeout" ]:
758+ retry_arg = getattr (self , f"retry_{ arg } )
759+ retry_arg [stage ] = getattr (
760+ self .args ,
761+ f"{ stage } { arg } ,
762+ None ,
763+ )
764764
765765 def raise_parser_error (self , message ):
766766 """
0 commit comments