Photo by Chris Ried on Unsplash

Django: How to write custom manage.py commands

Adrien Van Thong

--

The Django framework comes with a fairly extensive library of pre-packaged admin commands executed via the manage.py interface that can make managing your site and apps easier and more automatable. Did you know that this same framework also provides a mechanism to create your own brand-new customized admin commands?

Let’s explore how to do this and walk through the step-by-step process involved.

Set-up

Before we dive into the code, we need to establish a specific directory structure within our Django app for our custom commands:

my_app/
__init__.py
models.py
management/
__init__.py
commands/
__init__.py
my_custom_command.py
urls.py
views.py

The first change is a new management directory in any app which will have custom commands. Within this new directory we also need to create a sub-dir named commands which will have one python file for each custom command we’re creating.

It’s important that the file name matches the command that will invoke it. In the example above, my_custom_command.py will be run with the command: python manage.py my_custom_command. In the case of multiple custom commands, each command gets its own file inside the commands directory.

Next, let’s take a peek inside the my_custom_command.py file.

Defining our custom admin command

The my_custom_command.py file contains the customized logic which will run when the system administrator invokes the command: python manage.py my_custom_command

from django.core.management.base import BaseCommand

class Command(BaseCommand):
help = 'Customized admin command. Hi readers!'

def handle(self, *args, **options):
# Logic goes here!
self.stdout.write('This is a custom command. With custom output. Hello World!')

The my_custom_command.py file contains just a single class, always named Command and inheriting from BaseCommand class. The python file’s name is what determines the command name.

Within this class are two important components: the help attribute and the handle method.

The help attribute is a string which allows the developer to write the help text for the new custom command. This text is shown as the description next to the custom command when the administrator lists all available commands by running python manage.py without specifying a specific command.

The handle method is what contains all our customized code and where we will put all our logic. This can import other models and perform actions on them, or whatever other magic needs to happen.

Note that in this situation, we use the self.stdout.write method instead of print method to output text to the standard out. Conversely, self.stderr.write can be used for standard error when printing error messages and the like.

Here is what we see when we run our custom method:

$ python manage.py

Type 'manage.py help <subcommand>' for help on a specific subcommand.

Available subcommands:

[auth]
changepassword
createsuperuser

[contenttypes]
remove_stale_contenttypes

[django]
check
...
startapp
startproject
test
testserver

[my_app]
my_custom_command

$ python manage.py my_custom_command
This is a custom command. With custom output. Hello World!

For added pizzazz you can optionally color the output using the following style method:

self.stdout.write(self.style.SUCCESS('This is success text! Yay!'))

Other options are available such as ERROR, NOTICE, WARNING. Consult the Django docs for a complete list of custom color options when outputting text.

Lastly, the CommandError exception can be raised to indicate something went wrong and to terminate execution of our custom command. The Django framework will handle this exception gracefully for the user and print the exception text to standard error and exit the application with the specified returncode

Here is an example which puts it all together:

from django.core.management.base import BaseCommand, CommandError
from .models import Bar

class Command(BaseCommand):
help = 'Foos all enabled Bars in the system.'

def handle(self, *args, **options):
try:
Bar.objects.filter(enabled=True).update(foo=True)
except Exception as e:
raise CommandError(f'Something bad hapenned: {e}')
self.stdout.write(self.style.SUCCESS('Successfully Foo-ed all enabled Bars!'))

Passing in command-line arguments

There may come a time when we want to pass in command-line arguments to our new custom admin command. Fortunately, Django provides us a very familiar mechanism for doing so, leveraging the popular argparse python module.

To get started, we’ll need to override the add_arguments method, as such:

def add_arguments(self, parser):
parser.add_argument('id', type=int, help='ID of the record being updated')

In our handle method, we can access the arguments via the options parameter passed into the method. For example:

def handle(self, *args, **options):
obj_id = options.get('id')
try:
bar = Bar.objects.get(id=obj_id)
except Bar.DoesNotExist as e:
raise CommandError(f'Object with ID {obj_id} not found!')
bar.foo = True
bar.save()
self.stdout.write(self.style.SUCCESS('Bar has been Foo-ed!'))

When we run our new command here is what we see:

$ python manage.py my_custom_command 1
Bar has been Foo-ed!
$ python manage.py my_custom_command 2
CommandError: Object with ID 2 not found!

Sample Unit Test

As a good developer, we next want to unit test our new custom method. After all, any number of automated processes may be relying on our custom command to continue working as intended.

Django provides us with a very easy way to invoke our new custom admin command from the unit test framework (or anywhere else, for that matter) via the call_command method. This method takes in a string representing the admin command to run.

A sample unit test to make sure all our Bars were Fooed would look as follows:

from django.core.management import call_command
from django.test import TestCase
from .models import Bar

class CustomCommandTest(TestCase):
def test_all_bars_fooed(self):
Bar.objects.create(foo=False, enabled=False)
Bar.objects.create(foo=False, enabled=True)
call_command('my_custom_command')
self.assertEqual(1, Bar.objects.filter(foo=True).count())
self.assertEqual(1, Bar.objects.filter(foo=False).count())

As a best practice, I recommend abstracting away as much logic outside of the Command class into separate modules which can be tested individually and be easily re-used elsewhere.

That’s it! Django makes it very simple to customize our site for our sysadmins or other automated hooks to execute extra code all within the Django framework and app environment.

External Resources

--

--

No responses yet