Repository: sandialabs/staged-script

Issue:
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 had a `dict` where the keys are the stage names, mypy would be happy.

> **Note:**  This will be a breaking change.

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/sandialabs/staged-script/pull/40?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sandialabs) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 100.00%. Comparing base [(`1431129`)](https://app.codecov.io/gh/sandialabs/staged-script/commit/14311292744133fb3d72d93916b739b3b08c429f?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sandialabs) to head [(`dffc08c`)](https://app.codecov.io/gh/sandialabs/staged-script/commit/dffc08cecb8c711959fd5b1a0656213803a6e92c?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sandialabs).

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff            @@
##            master       #40   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files            2         2           
  Lines          180       186    +6     
  Branches        21        21           
=========================================
+ Hits           180       186    +6     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/sandialabs/staged-script/pull/40?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sandialabs).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sandialabs).

jmgate: Still need to check the Examples page in the docs to see if all is well.

Patch:
diff --git a/doc/source/examples.rst b/doc/source/examples.rst
index d3e8532..3568db4 100644
--- 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
index d63124a..219d7ea 100755
--- a/example/ex_1_removing_the_retry_arguments.py
+++ b/example/ex_1_removing_the_retry_arguments.py
@@ -31,12 +31,12 @@ class MyScript(StagedScript):
         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
index 4f70ed5..c067a56 100755
--- a/example/ex_2_running_certain_stages_by_default.py
+++ b/example/ex_2_running_certain_stages_by_default.py
@@ -33,12 +33,12 @@ class MyScript(StagedScript):
         )
         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
index 222dd15..1272102 100755
--- a/example/ex_3_adding_arguments.py
+++ b/example/ex_3_adding_arguments.py
@@ -36,12 +36,12 @@ class MyScript(StagedScript):
         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
index 2ac7422..d663b01 100755
--- a/example/ex_4_customizing_stage_behavior.py
+++ b/example/ex_4_customizing_stage_behavior.py
@@ -36,12 +36,12 @@ class MyScript(StagedScript):
         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
index b10f61a..df34051 100755
--- a/example/ex_5_customizing_individual_stages.py
+++ b/example/ex_5_customizing_individual_stages.py
@@ -36,12 +36,12 @@ class MyScript(StagedScript):
         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
index bf0eff8..1459692 100755
--- a/example/ex_6_creating_retryable_stages.py
+++ b/example/ex_6_creating_retryable_stages.py
@@ -61,12 +61,12 @@ class MyScript(StagedScript):
     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
index 4a0b173..5751f4a 100755
--- a/example/ex_7_customizing_the_summary.py
+++ b/example/ex_7_customizing_the_summary.py
@@ -63,12 +63,12 @@ class MyScript(StagedScript):
     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
index 7f4b9f1..15b8a22 100644
--- a/staged_script/staged_script.py
+++ b/staged_script/staged_script.py
@@ -16,6 +16,7 @@ import re
 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 @@ class StagedScript:
         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 @@ class StagedScript:
                 """
                 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 @@ class StagedScript:
             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 @@ class StagedScript:
 
         .. 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,7 +705,7 @@ class StagedScript:
                     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,
@@ -714,7 +713,7 @@ class StagedScript:
                     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,
@@ -722,7 +721,7 @@ class StagedScript:
                     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 @@ class StagedScript:
             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):
         """


Test patch:
diff --git a/test/test_staged_script.py b/test/test_staged_script.py
index f820729..b1f1ca8 100644
--- 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,