Writing custom management commands

November 14, 2008 Django, Programming, Python

The other night in the #django-dev IRC channel, Russ, Eric and I were talking about custom management commands for certain types of common but tedious tasks; Eric was discussing the possibility of a command for automatically generating a tests module in a Django application, since he’s our resident unit-testing freak, and I started toying with the idea of one to generate basic admin declarations for the models in an application.

So I sat down and spent a few minutes writing the code. I haven’t decided yet whether I want to clean it up into a proper patch against Django (which would require documentation and some other work), or just stuff it away somewhere to break out when I need it, but I thought the process would make an interesting write-up; although Brian and Michael have both written some tips for doing this, one can never have too many examples (and someday I plan to expand Django’s own documentation with more information, so think of this as a first draft).

First things first

Before writing any code, first it’s necessary to decide what, exactly, the command should do. In this case, I want a few things to happen:

The command should take one or more application names as its input.
If any of the named applications already have an admin.py file, or don’t define any models the admin could be enabled for, don’t do anything to those applications.
For an application which does have models and doesn’t have an admin.py file, create the admin.py file and populate it with bare-bones admin declarations for all the models in the app.
Leave a couple of useful comments in the generated admin.py file pointing people to the full admin documentation for customization.

Now I can put together the basic files needed for the command; since I may end up submitting it for inclusion in Django, I’ll stick it in the logical place: my local copy of django.contrib.admin. This involves creating a couple of directories and files inside django/contrib/admin:

A directory named management.
Inside it, an empty __init__.py (to turn the directory into a Python module) and a directory named commands.
Inside commands, files named __init__.py and createadmin.py (the name of the command will be taken from the name of the file it lives in, so in this case it will end up being manage.py createadmin).
Inside createadmin.py, create a class named Command which subclasses one of the base classes for management commands.

Anatomy of a management command

All Django management commands are classes descended from django.core.management.base.BaseCommand, which defines the basic infrastructure needed for a management command. But while you can directly subclass BaseCommand if you really need to, it’s often better to inherit from one of the other classes in django.core.management.base, since there are several useful subclasses of BaseCommand which make common patterns easier to handle:

AppCommand provides a simplified interface for writing a management command which takes a list of Django application names as arguments and does something with each of them. Built-in commands like sqlall make use of this.
LabelCommand does the same, but allows the arguments to be any strings you like, rather than only application names. This is used by, for example, startproject, which wants a name to give to the project (and, in fact, can create multiple projects in one go; startapp can similarly create multiple applications for you).
NoArgsCommand, which handles the much simpler case of a command which receives no arguments (aside from the standard django-admin.py and manage.py arguments like —verbosity and —help). The syncdb utility is a good example of this type of command.

Each of these subclasses of BaseCommand pares down the code you have to write yourself, to the point where writing a new command is a two-step process:

Create a class which inherits from the appropriate subclass of BaseCommand.
Override one method (the exact name varies according to the class you’re inheriting from), and have that method supply whatever custom logic you need, and a couple of attributes to provide information about the command.

In this case, I want to take a list of application names and do something with each of them, so AppCommand is the way to go; an AppCommand subclass simply has to define a method with the signature:

def handle_app(self, app, **options)

and AppCommand itself will handle the rest. This method will be called once for each application name passed on the command line; the argument app will be the models module of an application, and options will be a dictionary of any standard command-line arguments (e.g., —verbosity, which becomes the key verbosity in options) passed in.

Given this, I can already stub out the basics in createadmin.py:

from django.core.management.base import AppCommand

class Command(AppCommand):
    def handle_app(self, app, **options):
        pass

So far it doesn’t do anything, but that’ll change in a moment. Before moving on, there are two attributes which should be set to provide useful information to users:

from django.core.management.base import AppCommand

class Command(AppCommand):
    help = "Creates a basic admin.py file for the given app name(s)."

    args = "[appname ...]"

    def handle_app(self, app, **options):
        pass

The help attribute is, appropriately enough, the command’s brief description, which will be displayed when someone runs manage.py —help, and args is a description of the command-line arguments it expects (in this case, a list of application names).

Beginnings of the command’s logic

Let’s consider a simple example that pops up all over in Django tutorials and documentation: a weblog application, which we’ll think of as being named weblog, containing two models named Entry and Link. A minimal but customizable admin.py file for it would look like this:

from django.contrib import admin
from weblog.models import Entry, Link

class EntryAdmin(admin.ModelAdmin):
    pass


class LinkAdmin(admin.ModelAdmin):
    pass

admin.site.register(Entry, EntryAdmin)
admin.site.register(Link, LinkAdmin)

If there’s not going to be any customization, it could be simplified a bit, but I’d rather have a starting point that I can edit immediately than something I’ll have to change before I can work with it. So that output is going to be the goal.

Now, this means I need to generate some Python code programmatically; there are a few ways I could do this, but since Django has a bundled template system that can generate plain-text files of any type you like, I’ll just use a Django template to do this. Normally a Django application uses template loaders to find template files and compile them into Django Template objects, but there’s no requirement that you do so; you can create a Template object from any Python string. So I’ll just add a string containing Django template markup, and have my command compile it and use it to produce the output. And now the createadmin.py file looks like this:

from django.core.management.base import AppCommand

ADMIN_FILE_TEMPLATE = """from django.contrib import admin
from {{ app }} import {{ models|join:", " }}

{% for model in models %}
class {{ model }}Admin(admin.ModelAdmin):
    pass

{% endfor %}

{% for model in models %}admin.site.register({{ model }}, {{ model }}Admin)
{% endfor %}
"""


class Command(AppCommand):
    help = "Creates a basic admin.py file for the given app name(s)."

    args = "[appname ...]"

    def handle_app(self, app, **options):
        pass

Notice that it’s a triple-quoted string, which allows it to span across multiple lines. And it’s going to expect two variables: app, which will be the Python import path of the application’s models module (we’ll see in a moment how to get this from the app argument passed to the command), and models, which will be a list of the names of the model classes.

And now I can start filling in the command’s logic:

def handle_app(self, app, **options):
    from django import template
    from django.db.models import get_models

    models = get_models(app)

First of all, I’ll need the Django template module, so that’s imported. In order to find out which model classes are defined in the application, I’ll also need to use one of Django’s model-loading functions: get_models(), which takes an application’s models module as an argument and returns the model classes defined in it.

Next I can compile a Template object from the string I’ve provided:

t = template.Template(ADMIN_FILE_TEMPLATE)

And now I’ll need a Context to render it with, containing the two variables the template expects. The variable app should be the Python import path of the models module for the application (e.g., weblog.models), and luckily the argument app passed in to handle_app() is that module, so I can just use the __name__ attribute on it to get the necessary path (__name__ is the name of the module, in dotted import-path notation). And the variable models needs to be a list of the names of the model classes. For ease of access to this information, Django stores a usable name as part of the model’s options; given a Django model class in a variable named, say, m, then m._meta.object_name will be the class name. So here’s the context:

c = template.Context({ 'app': app.__name__, 'models': [m._meta.object_name for m in models] })

Now all that’s left for a first draft is to render the template with this context, and dump the output into a file named admin.py inside the application. This means I need to figure out the actual location, on disk, of the application; fortunately this is easy. The app argument to handle_app(), since it’s a Python module object, has an attribute named __file__ containing the path, on disk, to the file where that module is defined (so it’ll be something like “/Users/jbennett/dev/python/weblog/models.py”), and functions from Python’s standard os.path module will let me split that up to get the directory portion of that path (“/Users/jbennett/dev/python/weblog/”) and then tack a new filename onto the end of it:

admin_filename = os.path.join(os.path.dirname(app.__file__), 'admin.py')

Here’s how it works:

os.path.dirname takes a file path and returns only the path to the directory where that file lives.
os.path.join takes two or more bits of a file path and joins them together (following the operating system’s rules for this) to produce a single final path.

So if we’re looking at a weblog application whose models module is “/Users/jbennett/dev/python/weblog/models.py”, the code above will split off the directory portion, then join admin.py onto the end, to produce the path “/Users/jbennett/dev/python/weblog/admin.py”, which is precisely where the admin declarations need to go.

Then it’s just a matter of rendering the template and writing the output to the file:

admin_file = open(admin_filename, 'w')
admin_file.write(t.render(c))
admin_file.close()

And here’s the full code up to this point:

from django.core.management.base import AppCommand

ADMIN_FILE_TEMPLATE = """from django.contrib import admin
from {{ app }} import {{ models|join:", " }}

{% for model in models %}
class {{ model }}Admin(admin.ModelAdmin):
    pass

{% endfor %}

{% for model in models %}admin.site.register({{ model }}, {{ model }}Admin)
{% endfor %}
"""


class Command(AppCommand):
    help = "Creates a basic admin.py file for the given app name(s)."

    args = "[appname ...]"

    def handle_app(self, app, **options):
        import os.path
        from django import template
        from django.db.models import get_models

        models = get_models(app)
        t = template.Template(ADMIN_FILE_TEMPLATE)
        c = template.Context({ 'app': app.__name__, 'models': [m._meta.object_name for m in models] })

        admin_filename = os.path.join(os.path.dirname(app.__file__), 'admin.py')
        admin_file = open(admin_filename, 'w')
        admin_file.write(t.render(c))
        admin_file.close()

Improving the command

Of course, this is still missing a few things:

It tries to write an admin.py file regardless of whether there are any models in the application.
It will overwrite any existing admin.py file already defined in the application.
It doesn’t let you know what it’s doing.
It doesn’t provide any pointers on how to expand the admin definitions.

Dealing with the first problem is easy; we can simply look at the list of models returned by get_models(), and if it’s empty not bother doing anything:

if not models:
    return

The second problem is a simple matter of looking to see whether there’s already an admin.py file in the application, and another function in os.path can handle this: os.path.exists() will tell you whether the path you give it actually exists on disk. So all that’s needed is moving the admin.py generation into an if block:

if not os.path.exists(admin_filename):
    t = template.Template(ADMIN_FILE_TEMPLATE)
    c = template.Context({ 'app': app.__name__, 'models': [m._meta.object_name for m in models] })

    admin_file = open(admin_filename, 'w')
    admin_file.write(t.render(c))
    admin_file.close()

When os.path.exists returns True, this block of code won’t do anything, and the existing admin.py won’t be overwritten.

Providing feedback on what’s happening is similarly easy; all management commands allow the argument —verbosity to be passed in, specifying how much output should be printed during the command’s execution. As a general rule:

If verbosity is 0, print nothing.
If verbosity is 1 or greater, print basic information.
If verbosity is 2 or greater, print as much information as possible (typically for debugging).

The verbosity argument will end up inside the options argument if it’s specified, and it’s easy enough to read it out of there and supply a default (a value of 1 is usually good for this) in case it wasn’t specified:

verbosity = options.get('verbosity', 1)

Since the application’s name will be needed for printing messages about what’s going on, it can be pulled out of the app argument; this will, again, be the Python module where the application’s models are defined, and its __name__ attribute is the full dotted Python path of the module (e.g., weblog.models). So splitting it on the dots and pulling out the next-to-last item gets the application name:

app_name = app.__name__.split('.')[-2]

Then some basic information about what’s going on:

if verbosity > 0:
    print "Handling app '%s'" % app_name

Then a few other helpful messages can be sprinkled in; for example, when there are no models define in the application, some debugging output is useful:

if not models:
    if verbosity > 1:
        print "Skipping app '%s' : no models defined in this app" % app_name
    return

And similar messages can be added for the cases where an existing admin.py file is found.

Finally, some helpful comments can be added inside the template, to direct people to the documentation and explain how to use the file:

ADMIN_FILE_TEMPLATE = """from django.contrib import admin
from {{ app }} import {{ models|join:", " }}


# The following classes define the admin interface for your models.
# See http://docs.djangoproject.com/en/dev/ref/contrib/admin/ for
# a full list of the options you can use in these classes.
{% for model in models %}
class {{ model }}Admin(admin.ModelAdmin):
    pass

{% endfor %}
# Each of these lines registers the admin interface for one model. If
# you don't want the admin interface for a particular model, remove
# the line which registers it.
{% for model in models %}admin.site.register({{ model }}, {{ model }}Admin)
{% endfor %}
"""

And that’s that. The final, full command looks like this:

from django.core.management.base import AppCommand


ADMIN_FILE_TEMPLATE = """from django.contrib import admin
from {{ app }} import {{ models|join:", " }}


# The following classes define the admin interface for your models.
# See http://docs.djangoproject.com/en/dev/ref/contrib/admin/ for
# a full list of the options you can use in these classes.
{% for model in models %}
class {{ model }}Admin(admin.ModelAdmin):
    pass

{% endfor %}
# Each of these lines registers the admin interface for one model. If
# you don't want the admin interface for a particular model, remove
# the line which registers it.
{% for model in models %}admin.site.register({{ model }}, {{ model }}Admin)
{% endfor %}
"""


class Command(AppCommand):
    help = "Creates a basic admin.py file for the given app name(s)."

    args = '[appname ...]'

    def handle_app(self, app, **options):
        import os.path
        from django import template
        from django.db.models import get_models
        
        verbosity = options.get('verbosity', 1)
        app_name = app.__name__.split('.')[-2]
        if verbosity > 0:
           print "Handling app '%s'" % app_name
        
        models = get_models(app)
        
        if not models:
            if verbosity > 1:
                print "Skipping app '%s': no models defined in this app" % app_name
            return
        admin_filename = os.path.join(os.path.dirname(app.__file__), 'admin.py')
        if not os.path.exists(admin_filename):
            t = template.Template(ADMIN_FILE_TEMPLATE)
            c = template.Context({ 'app': app.__name__, 'models': [m._meta.object_name for m in models] })

            admin_file = open(admin_filename, 'w')
            admin_file.write(t.render(c))
            admin_file.close()
        else:
            if verbosity > 1:
                print "Skipping app '%s': admin.py file already exists" % app_name

This implements everything I wanted originally, so now I can simply use manage.py createadmin someappname someotherappname to start generating admin.py files for new applications (assuming I have django.contrib.admin in the INSTALLED_APPS of whatever project I’m working with; custom management commands are only loaded from applications in INSTALLED_APPS.

And that’s a wrap

Custom management commands take a moment or two to wrap your head around, but once you get them they’re a very powerful way to automate or simplify common Django-oriented tasks; lately I’ve started using them to provide an easier way to implement things which need to run in cron jobs (since manage.py takes care of setting DJANGO_SETTINGS_MODULE for you), for example in this custom command for django-registration, but they’re useful for all sorts of things and not terribly hard to write.