Skip to content

Commit

Permalink
refactor!: Retry argument attributes
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
jmgate committed Jul 3, 2024
1 parent bf9017a commit 730e6cb
Show file tree
Hide file tree
Showing 10 changed files with 84 additions and 92 deletions.
8 changes: 0 additions & 8 deletions doc/source/examples.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
12 changes: 6 additions & 6 deletions example/ex_1_removing_the_retry_arguments.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
12 changes: 6 additions & 6 deletions example/ex_2_running_certain_stages_by_default.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
12 changes: 6 additions & 6 deletions example/ex_3_adding_arguments.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
12 changes: 6 additions & 6 deletions example/ex_4_customizing_stage_behavior.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
12 changes: 6 additions & 6 deletions example/ex_5_customizing_individual_stages.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
12 changes: 6 additions & 6 deletions example/ex_6_creating_retryable_stages.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
12 changes: 6 additions & 6 deletions example/ex_7_customizing_the_summary.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
82 changes: 41 additions & 41 deletions staged_script/staged_script.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
import shlex
import subprocess
from argparse import (
Action,
ArgumentDefaultsHelpFormatter,
ArgumentParser,
Namespace,
Expand Down Expand Up @@ -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
Expand All @@ -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__(
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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"]
)
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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):
"""
Expand Down
2 changes: 1 addition & 1 deletion test/test_staged_script.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down

0 comments on commit 730e6cb

Please sign in to comment.