refactor!: Retry argument attributes

Mypy doesn't like that the retry argument attributes are defined dynamically when a class is instantiated, so it'll complain that they don't exist when a subclass developer is trying to modify them when overriding/extending the parser. If instead of having separate attributes per registered stage, we just have a `dict`` where the keys are the stage names, mypy is happy. Note that this is a breaking change, as it changes how subclass developers customize the parser.
sandialabs · Jul 3, 2024 · 730e6cb · 730e6cb
1 parent bf9017a
commit 730e6cb
Show file tree

Hide file tree

Showing 10 changed files with 84 additions and 92 deletions.
diff --git a/doc/source/examples.rst b/doc/source/examples.rst
@@ -76,14 +76,6 @@ by adding the following to the ``MyScript`` class:
    :lines: 28-40
    :caption: ``example/ex_1_removing_the_retry_arguments.py``
 
-.. note::
-
-   An upcoming release will refactor the retry argument attributes so
-   `mypy`_ will be happy with them.  For now, just use the ``type:
-   ignore[attr-defined]`` comments.
-
-.. _mypy:  https://mypy-lang.org/
-
 Now when we look at the ``--help`` text, we see:
 
 .. command-output:: python3 ../../example/ex_1_removing_the_retry_arguments.py --help

diff --git a/example/ex_1_removing_the_retry_arguments.py b/example/ex_1_removing_the_retry_arguments.py
@@ -31,12 +31,12 @@ def parser(self) -> ArgumentParser:
         my_parser.description = "Demonstrate removing the retry arguments."
         self.retry_arg_group.title = argparse.SUPPRESS
         self.retry_arg_group.description = argparse.SUPPRESS
-        self.hello_retry_attempts_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.hello_retry_delay_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.hello_retry_timeout_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.goodbye_retry_attempts_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.goodbye_retry_delay_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.goodbye_retry_timeout_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
+        self.retry_attempts_arg["hello"].help = argparse.SUPPRESS
+        self.retry_delay_arg["hello"].help = argparse.SUPPRESS
+        self.retry_timeout_arg["hello"].help = argparse.SUPPRESS
+        self.retry_attempts_arg["goodbye"].help = argparse.SUPPRESS
+        self.retry_delay_arg["goodbye"].help = argparse.SUPPRESS
+        self.retry_timeout_arg["goodbye"].help = argparse.SUPPRESS
         return my_parser
 
     def main(self, argv: List[str]) -> None:

diff --git a/example/ex_2_running_certain_stages_by_default.py b/example/ex_2_running_certain_stages_by_default.py
@@ -33,12 +33,12 @@ def parser(self) -> ArgumentParser:
         )
         self.retry_arg_group.title = argparse.SUPPRESS
         self.retry_arg_group.description = argparse.SUPPRESS
-        self.hello_retry_attempts_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.hello_retry_delay_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.hello_retry_timeout_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.goodbye_retry_attempts_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.goodbye_retry_delay_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.goodbye_retry_timeout_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
+        self.retry_attempts_arg["hello"].help = argparse.SUPPRESS
+        self.retry_delay_arg["hello"].help = argparse.SUPPRESS
+        self.retry_timeout_arg["hello"].help = argparse.SUPPRESS
+        self.retry_attempts_arg["goodbye"].help = argparse.SUPPRESS
+        self.retry_delay_arg["goodbye"].help = argparse.SUPPRESS
+        self.retry_timeout_arg["goodbye"].help = argparse.SUPPRESS
         my_parser.set_defaults(stage=list(self.stages))
         return my_parser
 

diff --git a/example/ex_3_adding_arguments.py b/example/ex_3_adding_arguments.py
@@ -36,12 +36,12 @@ def parser(self) -> ArgumentParser:
         my_parser.description = "Demonstrate adding arguments to the parser."
         self.retry_arg_group.title = argparse.SUPPRESS
         self.retry_arg_group.description = argparse.SUPPRESS
-        self.hello_retry_attempts_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.hello_retry_delay_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.hello_retry_timeout_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.goodbye_retry_attempts_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.goodbye_retry_delay_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.goodbye_retry_timeout_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
+        self.retry_attempts_arg["hello"].help = argparse.SUPPRESS
+        self.retry_delay_arg["hello"].help = argparse.SUPPRESS
+        self.retry_timeout_arg["hello"].help = argparse.SUPPRESS
+        self.retry_attempts_arg["goodbye"].help = argparse.SUPPRESS
+        self.retry_delay_arg["goodbye"].help = argparse.SUPPRESS
+        self.retry_timeout_arg["goodbye"].help = argparse.SUPPRESS
         my_parser.set_defaults(stage=list(self.stages))
         my_parser.add_argument(
             "--some-file",

diff --git a/example/ex_4_customizing_stage_behavior.py b/example/ex_4_customizing_stage_behavior.py
@@ -36,12 +36,12 @@ def parser(self) -> ArgumentParser:
         my_parser.description = "Demonstrate adding arguments to the parser."
         self.retry_arg_group.title = argparse.SUPPRESS
         self.retry_arg_group.description = argparse.SUPPRESS
-        self.hello_retry_attempts_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.hello_retry_delay_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.hello_retry_timeout_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.goodbye_retry_attempts_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.goodbye_retry_delay_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.goodbye_retry_timeout_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
+        self.retry_attempts_arg["hello"].help = argparse.SUPPRESS
+        self.retry_delay_arg["hello"].help = argparse.SUPPRESS
+        self.retry_timeout_arg["hello"].help = argparse.SUPPRESS
+        self.retry_attempts_arg["goodbye"].help = argparse.SUPPRESS
+        self.retry_delay_arg["goodbye"].help = argparse.SUPPRESS
+        self.retry_timeout_arg["goodbye"].help = argparse.SUPPRESS
         my_parser.set_defaults(stage=list(self.stages))
         my_parser.add_argument(
             "--some-file",

diff --git a/example/ex_5_customizing_individual_stages.py b/example/ex_5_customizing_individual_stages.py
@@ -36,12 +36,12 @@ def parser(self) -> ArgumentParser:
         my_parser.description = "Demonstrate adding arguments to the parser."
         self.retry_arg_group.title = argparse.SUPPRESS
         self.retry_arg_group.description = argparse.SUPPRESS
-        self.hello_retry_attempts_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.hello_retry_delay_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.hello_retry_timeout_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.goodbye_retry_attempts_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.goodbye_retry_delay_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.goodbye_retry_timeout_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
+        self.retry_attempts_arg["hello"].help = argparse.SUPPRESS
+        self.retry_delay_arg["hello"].help = argparse.SUPPRESS
+        self.retry_timeout_arg["hello"].help = argparse.SUPPRESS
+        self.retry_attempts_arg["goodbye"].help = argparse.SUPPRESS
+        self.retry_delay_arg["goodbye"].help = argparse.SUPPRESS
+        self.retry_timeout_arg["goodbye"].help = argparse.SUPPRESS
         my_parser.set_defaults(stage=list(self.stages))
         my_parser.add_argument(
             "--some-file",

diff --git a/example/ex_6_creating_retryable_stages.py b/example/ex_6_creating_retryable_stages.py
@@ -61,12 +61,12 @@ def say_goodbye(self) -> None:
     def parser(self) -> ArgumentParser:
         my_parser = super().parser
         my_parser.description = "Demonstrate adding arguments to the parser."
-        self.hello_retry_attempts_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.hello_retry_delay_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.hello_retry_timeout_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.goodbye_retry_attempts_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.goodbye_retry_delay_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.goodbye_retry_timeout_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
+        self.retry_attempts_arg["hello"].help = argparse.SUPPRESS
+        self.retry_delay_arg["hello"].help = argparse.SUPPRESS
+        self.retry_timeout_arg["hello"].help = argparse.SUPPRESS
+        self.retry_attempts_arg["goodbye"].help = argparse.SUPPRESS
+        self.retry_delay_arg["goodbye"].help = argparse.SUPPRESS
+        self.retry_timeout_arg["goodbye"].help = argparse.SUPPRESS
         my_parser.set_defaults(
             stage=list(self.stages),
             flaky_retry_attempts=5,

diff --git a/example/ex_7_customizing_the_summary.py b/example/ex_7_customizing_the_summary.py
@@ -63,12 +63,12 @@ def say_goodbye(self) -> None:
     def parser(self) -> ArgumentParser:
         my_parser = super().parser
         my_parser.description = "Demonstrate adding arguments to the parser."
-        self.hello_retry_attempts_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.hello_retry_delay_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.hello_retry_timeout_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.goodbye_retry_attempts_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.goodbye_retry_delay_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
-        self.goodbye_retry_timeout_arg.help = argparse.SUPPRESS  # type: ignore[attr-defined]
+        self.retry_attempts_arg["hello"].help = argparse.SUPPRESS
+        self.retry_delay_arg["hello"].help = argparse.SUPPRESS
+        self.retry_timeout_arg["hello"].help = argparse.SUPPRESS
+        self.retry_attempts_arg["goodbye"].help = argparse.SUPPRESS
+        self.retry_delay_arg["goodbye"].help = argparse.SUPPRESS
+        self.retry_timeout_arg["goodbye"].help = argparse.SUPPRESS
         my_parser.set_defaults(
             stage=list(self.stages),
             flaky_retry_attempts=5,

diff --git a/staged_script/staged_script.py b/staged_script/staged_script.py
@@ -16,6 +16,7 @@
 import shlex
 import subprocess
 from argparse import (
+    Action,
     ArgumentDefaultsHelpFormatter,
     ArgumentParser,
     Namespace,
@@ -74,6 +75,23 @@ class StagedScript:
         retry_arg_group (argparse._ArgumentGroup):  A container within
             the :class:`ArgumentParser` holding all the arguments
             associated with retrying stages.
+        retry_attempts (Dict[str, int]):  A mapping from stage names to
+            the number of times to attempt retrying each stage.
+        retry_attempts_arg (Dict[str, argparse.Action]):  The
+            corresponding arguments in the :class:`ArgumentParser`, so
+            subclass developers can modify them if needed.
+        retry_delay (Dict[str, float]):  A mapping from stage names to
+            how long to wait (in seconds) before attempting to retry a
+            stage.
+        retry_delay_arg (Dict[str, argparse.Action]):  The
+            corresponding arguments in the :class:`ArgumentParser`, so
+            subclass developers can modify them if needed.
+        retry_timeout (Dict[str, int]):  A mapping from stage names to
+            how long to wait (in seconds) before giving up on retrying
+            the stage.
+        retry_timeout_arg (Dict[str, argparse.Action]):  The
+            corresponding arguments in the :class:`ArgumentParser`, so
+            subclass developers can modify them if needed.
         script_name (str):  The name of the script (the
             :class:`StagedScript` subclass) being run.
         script_stem (str):  Same as :attr:`script_name`, but without the
@@ -90,28 +108,6 @@ class StagedScript:
             the user via the command line arguments.
         start_time (datetime):  The time at which this object was
             initialized.
-
-    Note that additional attributes are automatically generated for each
-    ``stage`` registered for a subclass object:
-
-    ``STAGE_NAME_retry_attempts`` (int)
-        The number of times to attempt retrying the ``STAGE_NAME``
-        stage.
-    ``STAGE_NAME_retry_attempts_arg`` (argparse.Action)
-        The corresponding argument in the :class:`ArgumentParser`, so
-        subclass developers can modify it if needed.
-    ``STAGE_NAME_retry_delay`` (float)
-        How long to wait (in seconds) before attempting to retry the
-        ``STAGE_NAME`` stage.
-    ``STAGE_NAME_retry_delay_arg`` (argparse.Action)
-        The corresponding argument in the :class:`ArgumentParser`, so
-        subclass developers can modify it if needed.
-    ``STAGE_NAME_retry_timeout`` (int)
-        How long to wait (in seconds) before giving up on retrying the
-        ``STAGE_NAME`` stage.
-    ``STAGE_NAME_retry_timeout_arg`` (argparse.Action)
-        The corresponding argument in the :class:`ArgumentParser`, so
-        subclass developers can modify it if needed.
     """
 
     def __init__(
@@ -155,6 +151,12 @@ def __init__(
         self.dry_run = False
         self.durations: List[StageDuration] = []
         self.print_commands = print_commands
+        self.retry_attempts: Dict[str, int] = {}
+        self.retry_attempts_arg: Dict[str, Action] = {}
+        self.retry_delay: Dict[str, float] = {}
+        self.retry_delay_arg: Dict[str, Action] = {}
+        self.retry_timeout: Dict[str, int] = {}
+        self.retry_timeout_arg: Dict[str, Action] = {}
         self.script_name = Path(__main__.__file__).name
         self.script_stem = Path(__main__.__file__).stem
         self.script_success = True
@@ -320,9 +322,9 @@ def wrapper(self, *args, **kwargs) -> None:
                 """
                 self.current_stage = stage_name
                 get_phase_method(self, "_run_pre_stage_actions")()
-                timeout = getattr(self, f"{stage_name}_retry_timeout")
-                attempts = getattr(self, f"{stage_name}_retry_attempts")
-                delay = getattr(self, f"{stage_name}_retry_delay")
+                timeout = self.retry_timeout[stage_name]
+                attempts = self.retry_attempts[stage_name]
+                delay = self.retry_delay[stage_name]
                 stop_after_timeout = stop_after_delay(timeout)
                 stop_after_max_attempts = stop_after_attempt(attempts + 1)
                 retry = Retrying(
@@ -590,10 +592,7 @@ def _handle_stage_retry_error_test(
             retry:  The :class:`Retrying` controller, which contains
                 information about the retrying that was done.
         """
-        retry_attempts = getattr(
-            self, f"{self.current_stage}_retry_attempts", 0
-        )
-        if retry_attempts > 0:
+        if self.retry_attempts[self.current_stage] > 0:
             stage_time = timedelta(
                 seconds=retry.statistics["delay_since_first_attempt"]
             )
@@ -660,9 +659,9 @@ def parser(self) -> ArgumentParser:
 
         .. code-block:: python
 
-            self.foo_retry_attempts_arg.help = argparse.SUPPRESS
-            self.foo_retry_delay_arg.help = argparse.SUPPRESS
-            self.foo_retry_timeout_arg.help = argparse.SUPPRESS
+            self.retry_attempts_arg["foo"].help = argparse.SUPPRESS
+            self.retry_delay_arg["foo"].help = argparse.SUPPRESS
+            self.retry_timeout_arg["foo"].help = argparse.SUPPRESS
 
         And if you want to remove the title for the retry group
         altogether, you can do so with:
@@ -706,23 +705,23 @@ def parser(self) -> ArgumentParser:
                     type=int,
                     help=f"How many times to retry the {stage!r} stage.",
                 )
-                setattr(self, f"{stage}_retry_attempts_arg", retry_attempts)
+                self.retry_attempts_arg[stage] = retry_attempts
                 retry_delay = self.retry_arg_group.add_argument(
                     f"--{stage}-retry-delay",
                     default=0,
                     type=float,
                     help="How long to wait (in seconds) before retrying the "
                     f"{stage!r} stage.",
                 )
-                setattr(self, f"{stage}_retry_delay_arg", retry_delay)
+                self.retry_delay_arg[stage] = retry_delay
                 retry_timeout = self.retry_arg_group.add_argument(
                     f"--{stage}-retry-timeout",
                     default=60,
                     type=int,
                     help="How long to wait (in seconds) before giving up on "
                     f"retrying the {stage!r} stage.",
                 )
-                setattr(self, f"{stage}_retry_timeout_arg", retry_timeout)
+                self.retry_timeout_arg[stage] = retry_timeout
         return my_parser
 
     def parse_args(self, argv: List[str]) -> None:
@@ -755,12 +754,13 @@ def parse_args(self, argv: List[str]) -> None:
             set(self.args.stage) if self.args.stage is not None else set()
         )
         for stage in self.stages:
-            for retry_arg in [
-                f"{stage}_retry_attempts",
-                f"{stage}_retry_delay",
-                f"{stage}_retry_timeout",
-            ]:
-                setattr(self, retry_arg, getattr(self.args, retry_arg, None))
+            for arg in ["attempts", "delay", "timeout"]:
+                retry_arg = getattr(self, f"retry_{arg}")
+                retry_arg[stage] = getattr(
+                    self.args,
+                    f"{stage}_retry_{arg}",
+                    None,
+                )
 
     def raise_parser_error(self, message):
         """

diff --git a/test/test_staged_script.py b/test/test_staged_script.py
@@ -98,7 +98,7 @@ def test__handle_stage_retry_error(
 ) -> None:
     """Test the :func:`_handle_stage_retry_error` method."""
     script.current_stage = "test"
-    script.test_retry_attempts = retry_attempts  # type: ignore[attr-defined]
+    script.retry_attempts["test"] = retry_attempts
     retry = mock_Retrying()
     retry.statistics = {
         "delay_since_first_attempt": 1234,