Django: create sub-commands within a management command

A tool of many sub-tools.

argparse, the standard library module that Django uses for parsing command line options, supports sub-commands. These are pretty neat for providing an expansive API without hundreds of individual commands. Here’s an example of using sub-commands in a Django management command:

from django.core.management.base import BaseCommand


class Command(BaseCommand):
    def add_arguments(self, parser):
        subparsers = parser.add_subparsers(
            title="sub-commands",
            required=True,
        )

        recharge_parser = subparsers.add_parser(
            "recharge",
            help="Recharge the laser.",
        )
        recharge_parser.set_defaults(method=self.recharge)

        shine_parser = subparsers.add_parser(
            "shine",
            help="Shine the laser.",
        )
        shine_parser.add_argument(
            "--bright", action="store_true", help="Make it lighter."
        )
        shine_parser.set_defaults(method=self.shine)

    def handle(self, *args, method, **options):
        method(*args, **options)

    def recharge(self, *args, **options):
        self.stdout.write("Recharging the laser...")

    def shine(self, *args, bright, **options):
        if bright:
            self.stdout.write("✨✨✨ Shine ✨✨✨")
        else:
            self.stdout.write("✨ Shine ✨")

The add_arguments() method sets up two sub-parsers for different sub-commands: recharge and shine. They use a pattern recommended by the argparse add_subparsers() documentation, storing a function to call for the parser with set_defaults(). The second sub-parser for shine takes an optional --bright argument as well.

The parser.add_subparsers() call sets required=True, which makes argparse show the command help and exit if no sub-command is provided. This means handle() is only called when a sub-command is selected, so its function signature can directly reference the method keyword argument. handle() only route to the correct method by calling method with the remaining arguments. (It could do some stuff common to all sub-commands, too.)

The sub-command methods are then defined to take their individual actions. The arguments passed to them come from their specific sub-parsers.

Invoking the sub-commands works as you might hope:

$ ./manage.py laser recharge
Recharging the laser...

$ ./manage.py laser shine
✨ Shine ✨

$ ./manage.py laser shine --bright
✨✨✨ Shine ✨✨✨

Running the command without a sub-command shows a “arguments are required” message from argparse:

$ ./manage.py laser
usage: manage.py laser [-h] [--version] [-v {0,1,2,3}] [--settings SETTINGS] [--pythonpath PYTHONPATH] [--traceback] [--no-color] [--force-color] [--skip-checks] {recharge,shine} ...
manage.py laser: error: the following arguments are required: {recharge,shine}

The help message shows the sub-commands, albeit after the full list of default options from Django:

$ ./manage.py laser --help
usage: manage.py laser [-h] [--version] [-v {0,1,2,3}] [--settings SETTINGS] [--pythonpath PYTHONPATH] [--traceback] [--no-color] [--force-color] [--skip-checks] {recharge,shine} ...

options:
  -h, --help            show this help message and exit
  --version             Show program's version number and exit.
  -v {0,1,2,3}, --verbosity {0,1,2,3}
  ...
  --skip-checks         Skip system checks.

sub-commands:
  {recharge,shine}
    recharge            Recharge the laser.
    shine               Shine the laser.

Sub-commands also have their own help pages:

$ ./manage.py laser shine --help
usage: manage.py laser shine [-h] [--bright]

options:
  -h, --help  show this help message and exit
  --bright    Make it lighter.

Testing sub-commands

The sub-commands can be tested like normal command arguments, using Django’s call_command(), per my previous post.

from io import StringIO

from django.core.management import call_command
from django.test import SimpleTestCase


class LaserTests(SimpleTestCase):
    def call_command(self, *args, **kwargs):
        out = StringIO()
        err = StringIO()
        call_command(
            "laser",
            *args,
            stdout=out,
            stderr=err,
            **kwargs,
        )
        return out.getvalue(), err.getvalue()

    def test_recharge(self):
        out, err = self.call_command("recharge")

        self.assertEqual(out, "Recharging the laser...\n")
        self.assertEqual(err, "")

    def test_shine(self):
        out, err = self.call_command("shine")

        self.assertEqual(out, "✨ Shine ✨\n")
        self.assertEqual(err, "")

    def test_shine_bright(self):
        out, err = self.call_command("shine", "--bright")

        self.assertEqual(out, "✨✨✨ Shine ✨✨✨\n")
        self.assertEqual(err, "")

Musing on Django’s built-in commands

Using sub-commands could simplify some of Django’s built-in commands. In particular, I am thinking of the large family of migration commands: makemigrations, migrate, optimizemigration, and so on.

I would find it more intuitive to use commands like ./manage.py migrations make and ./manage.py migrations run. It would also make the lesser-used commands a bit more discoverable through the CLI help.

But it would probably be too much upheaval to remove the existing commands, so this would have to be a new CLI. On balance, it’s probably not worth introducing a second “way of doing things”. Ah well, maybe if Django adds another group of commands in the future. (Maybe for background tasks?).

Fin

May your commands and sub-commands be well-tested and performant,

—Adam


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


Subscribe via RSS, Twitter, Mastodon, or email:

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

Related posts:

Tags: