A Python Script Template with Sub-commands (and Type Hints)

2021-10-15 Now there are two of them!

Earlier this week I shared my Python script template. Here’s an extended version with sub-command support, and an example script.

The Template

Voilà:

#!/usr/bin/env python
from __future__ import annotations

import argparse
from collections.abc import Sequence


def main(argv: Sequence[str] | None = None) -> int:
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="command", required=True)

    # Add sub-commands
    subcommand1_parser = subparsers.add_parser("subcommand1", help="...")
    subcommand1_parser.add_argument("string")

    args = parser.parse_args(argv)

    if args.command == "subcommand1":
        return subcommand1(args.string)
    else:
        # Unreachable
        raise NotImplementedError(
            f"Command {args.command} does not exist.",
        )


def subcommand1(string: str) -> int:
    # Implement behaviour

    return 0


if __name__ == "__main__":
    raise SystemExit(main())

This works with Python 3.7+.

Here’s what’s different from vanilla template:

add_subparsers() configures our parser to use sub-commands. We tell it to store the command name for later comparison, and make naming a sub-command required.
We add parsers for each sub-command with subparsers.add_parser(), which returns an ArgumentParser. We can use this to add arguments in the normal way. (...or even configure sub-sub-commands by calling its add_subparsers().)
After argument parsing, we use if to switch on the sub-command name in args.command, and call the corresponding sub-command functions. We deconstruct args when calling, so that our sub-command functions can take precise types.

Switching with if better than the set_defaults(func=...) method mentioned first in the argparse docs. Using if keeps control flow clear, and it means our sub-command functions can take individual arguments with types, rather than args.
Our else block raises an error for unrecognized sub-commands. This is not normally reachable - argparse already exits with a message for unknown sub-commands:
```
$ ./example.py subcommand2
usage: example.py [-h] {subcommand1} ...
example.py: error: argument command: invalid choice: 'subcommand2' (choose from 'subcommand1')
```
But if we add a sub-command parser and miss its corresponding if block, the error will become possible. So it’s best to keep it around.

To support use without a sub-command, remove required=True from the add_subparsers() call, and add behaviour in this else block.

Por Exemplo

Here we have two sub-commands, debug and hello:

#!/usr/bin/env python
from __future__ import annotations

import argparse
import sys
from collections.abc import Sequence


def main(argv: Sequence[str] | None = None) -> int:
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="command", required=True)

    subparsers.add_parser("debug", help="Show debug information.")

    hello_parser = subparsers.add_parser("hello", help="Say hello.")
    hello_parser.add_argument("name", help="Who to greet.")

    args = parser.parse_args(argv)

    if args.command == "debug":
        return debug()
    elif args.command == "hello":
        return hello(name=args.name)
    else:
        raise NotImplementedError(
            f"Command {args.command} does not exist.",
        )


def debug() -> int:
    print(f"Python version {sys.version}")
    return 0


def hello(name: str) -> int:
    print(f"Hello {name}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

Cool... let’s try it out.

We can ask for help and argparse obliges:

$ ./example.py -h
usage: example.py [-h] {debug,hello} ...

positional arguments:
  {debug,hello}
    debug        Show debug information.
    hello        Say hello.

optional arguments:
  -h, --help     show this help message and exit

We can also get help for sub-commands:

$ ./example.py hello --help
usage: example.py hello [-h] name

positional arguments:
  name        Who to greet.

optional arguments:
  -h, --help  show this help message and exit

The debug sub-command is straightforward:

$ ./example.py debug
Python version 3.9.7 (default, Sep  1 2021, 11:36:38)
[Clang 12.0.5 (clang-1205.0.22.11)]

And hello works as expected:

$ ./example.py hello reader
Hello reader

Very cool!

Fin

Go and create some great tools,

—Adam

😸😸😸 Check out my new book on using GitHub effectively, Boost Your GitHub DX! 😸😸😸

One summary email a week, no spam, I pinky promise.

Related posts:

A Python Script Template, with and without Type Hints and Async

Tags: mypy, python