Reference

The staged_script module.

Provides the StagedScript base class to extend when creating your own staged scripts, along with some helpers.

The StagedScript Class

class staged_script.staged_script.StagedScript(stages, *, console_force_terminal=None, console_log_path=True, print_commands=True)[source]

Bases: object

The base class for all staged scripts.

This serves as a base class for any Python scripts intended to drive a series of commands in the underlying shell. It’s built around the fundamental concept of a “stage”, in that such scripts are typically broken down into stages, each of which consists of a handful of commands. Certain actions are performed at the beginning or end of any given stage, and it’s also possible for stages to be skipped or retried. The class also provides a means of producing a script execution summary to show the user exactly what was done, for the sake of replicability and easing debugging.

args

The parsed command line arguments for the script.

Type:

Namespace

commands_executed

The commands that were executed in the shell.

Type:

List[str]

console

Used to print rich text to the console.

Type:

Console

current_stage

The name of the stage being run.

Type:

str

dry_run

If True, don’t actually run the command that would be executed in the shell; instead just print it out.

Type:

bool

durations

A mapping from stage names to how long it took for each to run. This is implemented as a list of named tuples instead of as a dict to allow the flexibility for stages to be run multiple times.

Type:

List[StageDuration]

print_commands

Whether to print the commands executed immediately before executing them.

Type:

bool

retry_arg_group

A container within the ArgumentParser holding all the arguments associated with retrying stages.

Type:

argparse._ArgumentGroup

script_name

The name of the script (the StagedScript subclass) being run.

Type:

str

script_stem

Same as script_name, but without the extension.

Type:

str

script_success

Subclass developers can toggle this attribute to indicate whether the script has succeeded.

Type:

bool

stage_start_time

The time at which a stage began.

Type:

datetime

stages

The stages registered for an instantiation of a StagedScript subclass, which are used to automatically populate pieces of the ArgumentParser. This may be a subset of all the stages defined in the subclass.

Type:

Set[str]

stages_to_run

Which stages to run, as specified by the user via the command line arguments.

Type:

Set[str]

start_time

The time at which this object was initialized.

Type:

datetime

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 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 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 ArgumentParser, so subclass developers can modify it if needed.

Initialize a StagedScript object.

Parameters:
  • stages (Set[str]) – The set of stages to register for a StagedScript subclass, which are used to automatically populate pieces of the ArgumentParser. This may be a subset of all the stages defined in the subclass.

  • console_force_terminal (Optional[bool]) – Whether to force the console to behave like a terminal. None allows auto-detection.

  • console_log_path (bool) – Whether to print the location within a file that generated a line in the console log.

  • print_commands (bool) – Whether to print the commands executed immediately before executing them.

Note

If you override this constructor in a subclass—e.g., to create additional attributes, etc.—you’ll need to call this parent constructor with super().__init__(stages) and optionally pass in additional arguments.

Stage Definition

@staged_script.staged_script.StagedScript.stage(stage_name, heading)

Convert a method to a stage.

A decorator to take a function and convert it to a conceptual stage of a script. Each stage consists of the following phases:

Pre-Stage Actions

A series of commands to run before a stage technically begins. See _run_pre_stage_actions().

Begin-Stage Actions

A series of commands to run at the beginning of the stage. See _begin_stage().

Stage Body

The actual work of the stage itself, encapsulated in the decorated function. However, if a particular stage is not to be executed (if it’s not passed in with the --stage flag on the command line), then there’s also the concept of Skip-Stage Actions, which are a series of commands to run if a stage is to be skipped. See _skip_stage().

End-Stage Actions

A series of commands to run at the end of the stage, even if an exception was raised within the Stage Body. See _end_stage().

Post-Stage Actions

A series of commands to run after a stage has technically wrapped up. See _run_post_stage_actions().

If a subclass developer writes a function to be wrapped by this decorator such that it raises a RetryStage exception on certain failure conditions, then there are additional Prepare-to-Retry Actions, which are a series of commands to run before the next attempt at running the stage (see _prepare_to_retry_stage()). The Begin-Stage Actions, Stage Body, and End-Stage Actions are wrapped in this retry loop, while the Pre- and Post-Stage Actions are not.

If the retry specifications (see parser()) are exhausted and the wrapped function still raises a RetryStage exception, then there is a Retry Error Handler containing a series of commands to run when exiting the retry loop (see _handle_stage_retry_error()).

Parameters:
  • stage_name (str) – The name of the stage, which must consist of only lowercase letters. Note that stage names must be unique within a class that inherits from StagedScript.

  • heading (str) – A heading message to print indicating what will happen in the stage.

Return type:

Callable

Stage Customization

StagedScript._run_pre_stage_actions()[source]

Run a series of commands before a stage actually starts.

This is implemented here as a no-op, but subclass developers can override it if they wish to always perform certain actions before stages get underway (e.g., maybe you want to log some details on the state of the system, etc.).

Subclass developers can also customize the Pre-Stage Actions for an individual stage by defining a _run_pre_stage_actions_STAGE_NAME method, where STAGE_NAME should be replaced with the name of the stage (the first argument to the stage() decorator). This can be useful if, e.g., you wanted to ensure certain pre-conditions were met before attempting the stage, and erroring out appropriately if not.

Return type:

None

StagedScript._begin_stage(heading)[source]

Execute a series of commands at the beginning of every stage.

If subclass developers wish to extend the Begin-Stage Actions, they can do so with the following:

def _begin_stage(
    self,
    heading: str
) -> None:
    super()._begin_stage(heading)
    # Insert more actions here.

Alternatively you can override the default behavior entirely by omitting the super() line above.

In addition, you may wish to customize the Begin-Stage Actions for a particular stage, in which case you define a _begin_stage_STAGE_NAME method, where STAGE_NAME should be replaced with the name of the stage (the first argument to the stage() decorator), e.g.:

def _begin_stage_test(  # Particular to the 'test' stage.
    self,
    heading: str
) -> None:
    self._begin_stage(heading)  # Optional
    # Insert more actions here.
Parameters:

heading (str) – A heading message to print indicating what will happen in the stage.

Return type:

None

StagedScript._skip_stage()[source]

Execute a series of commands when skipping a stage.

If subclass developers wish to extend the Skip-Stage Actions, they can do so with the following:

def _skip_stage(self) -> None:
    super()._skip_stage()
    # Insert more actions here.

Alternatively you can override the default behavior entirely by omitting the super() line above.

In addition, you may wish to customize the Skip-Stage Actions for a particular stage, in which case you define a _skip_stage_STAGE_NAME method, where STAGE_NAME should be replaced with the name of the stage (the first argument to the stage() decorator), e.g.:

# Particular to the 'test' stage.
def _skip_stage_test(self) -> None:
    self._skip_stage()  # Optional
    # Insert more actions here.
Return type:

None

StagedScript._end_stage()[source]

Execute a series of commands at the end of every stage.

If subclass developers wish to extend the End-Stage Actions, they can do so with the following:

def _end_stage(self) -> None:
    super()._end_stage()
    # Insert more actions here.

Alternatively you can override the default behavior entirely by omitting the super() line above.

In addition, you may wish to customize the End-Stage Actions for a particular stage, in which case you define an _end_stage_STAGE_NAME method, where STAGE_NAME should be replaced with the name of the stage (the first argument to the stage() decorator), e.g.:

# Particular to the 'test' stage.
def _end_stage_test(self) -> None:
    self._end_stage()  # Optional
    # Insert more actions here.
Return type:

None

StagedScript._run_post_stage_actions()[source]

Run a series of commands after a stage wraps up.

This is implemented here as a no-op, but subclass developers can override it if they wish to always perform certain actions after stages complete and before execution moves on with the rest of the script (e.g., maybe you want to log some details on the state of the system, etc.).

Subclass developers can also customize the Post-Stage Actions for an individual stage by defining a _run_post_stage_actions_STAGE_NAME method, where STAGE_NAME should be replaced with the name of the stage (the first argument to the stage() decorator). This can be useful if, e.g., you wanted to ensure certain post-conditions were met before moving on, and erroring out appropriately if not.

Return type:

None

StagedScript._prepare_to_retry_stage(retry_state)[source]

Prepare to retry a stage.

This method will be executed after the End-Stage Actions of one attempt and before the Begin-Stage Actions of a following attempt.

If subclass developers wish to extend the Prepare-to-Retry Actions, they can do so with the following:

def _prepare_to_retry_stage(
    self,
    retry_state: RetryCallState
) -> None:
    super()._prepare_to_retry_stage(retry_state)
    # Insert more actions here.

Alternatively you can override the default behavior entirely by omitting the super() line above.

In addition, you may wish to customize the Prepare-to-Retry Actions for an individual stage by defining a _prepare_to_retry_stage_STAGE_NAME method, where STAGE_NAME should be replaced with the name of the stage (the first argument to the stage() decorator), e.g.:

# Particular to the 'test' stage.
def _prepare_to_retry_stage_test(
    self,
    retry_state: RetryCallState
) -> None:
    self._prepare_to_retry_stage(retry_state)  # Optional
    # Insert more actions here.

This can be useful if, e.g., you need to reset some state before running a particular stage again.

Parameters:

retry_state (RetryCallState) – Information regarding the retry operation in progress.

Return type:

None

StagedScript._handle_stage_retry_error(retry)[source]

Handle a stage retry error.

When a stage has exhausted the specified retry conditions, handle the thrown RetryStage appropriately.

If subclass developers wish to extend the Retry Error Handler, they can do so with the following:

def _handle_stage_retry_error(
    self,
    retry: Retrying
) -> None:
    super()._handle_stage_retry_error(retry)
    # Insert more actions here.

Alternatively you can override the default behavior entirely by omitting the super() line above.

In addition, you may wish to customize the Retry Error Handler for an individual stage by defining a _handle_stage_retry_error_STAGE_NAME method, where STAGE_NAME should be replaced with the name of the stage (the first argument to the stage() decorator), e.g.:

# Particular to the 'test' stage.
def _handle_stage_retry_error_test(
    self,
    retry: Retrying
) -> None:
    self._handle_stage_retry_error(retry)  # Optional
    # Insert more actions here.
Parameters:

retry (Retrying) – The Retrying controller, which contains information about the retrying that was done.

Return type:

None

Parsing Arguments

property StagedScript.parser: ArgumentParser

The base argument parser.

Create an ArgumentParser for all the arguments made available by this base class. This should be overridden in subclasses as follows:

@functools.cached_property
def parser(self) -> ArgumentParser:
    my_parser = super().parser
    my_parser.description = "INSERT DESCRIPTION HERE"
    my_parser.add_argument("--foo", ...)
    ...
    # Optional changes.
    my_parser.formatter_class = RawDescriptionHelpFormatter
    my_parser.set_defaults(stage=list(self.stages))
    return my_parser

The formatter class defaults to a combination of ArgumentDefaultsHelpFormatter and RawDescriptionHelpFormatter, but can optionally be set to whatever you like.

If you’d like to set the default value for the --stage argument, you can use set_defaults() (the example above defaults to all the stages registered when the subclass was instantiated.

Similarly, if you wish to override the default values for the retry arguments that are automatically provided for every stage, you can do so with, e.g.:

my_parser.set_defaults(
    foo_retry_attempts=5,
    foo_retry_delay=10,
    foo_retry_timeout=600
)

If you wish to suppress the stage retry arguments for your stages that don’t raise a RetryStage exception, you can do so with, e.g.:

self.foo_retry_attempts_arg.help = argparse.SUPPRESS
self.foo_retry_delay_arg.help = argparse.SUPPRESS
self.foo_retry_timeout_arg.help = argparse.SUPPRESS

And if you want to remove the title for the retry group altogether, you can do so with:

self.retry_arg_group.title = argparse.SUPPRESS
self.retry_arg_group.description = argparse.SUPPRESS
Returns:

The base argument parser.

StagedScript.parse_args(argv)[source]

Parse the command line arguments.

Parse the command line arguments supplied by this base class. The recommendation is to save any arguments, perhaps with some transformations, as attributes of your subclass for ease of access throughout your subclass. This should be overridden in subclasses as follows:

def parse_args(self, argv: List[str]) -> None:
    super().parse_args(argv)
    # Parse additional arguments and store as attributes.
    self.foo = self.args.foo
    ...
    # Ensure supplied arguments are valid, etc.
    ...
Parameters:

argv (List[str]) – The command line arguments used when running this file as a script.

Return type:

None

StagedScript.raise_parser_error(message)[source]

Raise a parser error.

Exit the script with a message indicating what went wrong when parsing the command line arguments. To be used by subclass developers in the midst of parse_args().

Parameters:

message – What went wrong.

Raises:

SystemExit – To indicate the problem and stop script execution.

Running Commands

StagedScript.run(command, *, pretty_print=False, print_command=None, **kwargs)[source]

Run a command in the underlying shell.

Parameters:
  • command (str) – The command to be executed.

  • pretty_print (bool) – Whether the command should be “pretty-printed” when storing it in the list of commands executed.

  • print_command (Optional[bool]) – Whether to print the command executed immediately before executing it. If specified, this overrides the print_commands specified when instantiating the class.

  • kwargs – Additional keyword arguments to pass on to subprocess.run().

Return type:

CompletedProcess

Returns:

The result from calling subprocess.run().

StagedScript.pretty_print_command(command, indent=4)[source]

Pretty-print a command.

Take a command executed in the shell and pretty-print it by inserting newlines where appropriate:

  • A long-style flag with an argument is placed on its own line, e.g., --foo bar.

  • A long-style flag without an argument is placed on its own line.

  • Any other arguments (e.g., short-style flags, positional arguments) are placed on their own lines.

Parameters:
  • command (str) – The input command.

  • indent (int) – How many spaces to indent each subsequent line.

Return type:

str

Returns:

The reformatted command.

Script Execution Summary

StagedScript.print_script_execution_summary(extra_sections=None)[source]

Print a summary of everything that was done by the script.

Parameters:

extra_sections (Optional[Dict[str, str]]) – Any sections to add to the summary in addition to what’s automatically supplied by this base class. The keys are section headings, and values are the associated details. If you provide section headings identical to any supplied by this base class, the corresponding details will override that which is provided by default.

Return type:

None

Note: If you wish to override this in a subclass for the sake of providing additional sections every time it’s called, you can do something like the following:

def print_script_execution_summary(
    self,
    extra_sections: Optional[Dict[str, str]] = None
) -> None:
    extras = {"Additional section": "With some details."}
    if extra_sections is not None:
        extras.update(extra_sections)
    super().print_script_execution_summary(
        extra_sections=extras
    )

Additional Methods

StagedScript.print_heading(message, *, color='cyan')[source]

Print a heading.

To indicate at a high level what the script is doing.

Parameters:
  • message (str) – The message to print.

  • color (str) – What color to print the message in.

Return type:

None

StagedScript.print_dry_run_message(message, *, indent=0)[source]

Print a dry-run message.

Print a message indicating that something is happening due to the script running in dry-run mode.

Parameters:
  • message (str) – The message to print.

  • indent (int) – How many spaces by which to indent that which is printed.

Return type:

None

Helper Classes

class staged_script.staged_script.StageDuration(stage: str, duration: timedelta)[source]

Bases: NamedTuple

The duration of a stage.

Define a mapping from stage names to how long they took to run. See StagedScript.durations for details.

Create new instance of StageDuration(stage, duration)

duration: timedelta

Alias for field number 1

stage: str

Alias for field number 0

class staged_script.staged_script.RetryStage[source]

Bases: TryAgain

Trigger the retry of a stage.

Define an exception to be raised by subclass developers to indicate that a stage should be retried.

class staged_script.staged_script.HelpFormatter(prog, indent_increment=2, max_help_position=24, width=None)[source]

Bases: ArgumentDefaultsHelpFormatter, RawDescriptionHelpFormatter

The help formatter for the argument parser.

Define a formatter class to be used by the argument parser that both treats the description as raw text (doesn’t do any automatic formatting) and shows default values of arguments.