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

Learn how to make your tests run quickly in my book Speed Up Your Django Tests.

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