Easy Django choices done right

In a old, but still relevant post James Bennet argues for a way of handling the choice tuples used by Django. One of the reasons I haven’t always followed that advice was that it was just a lot of extra mundane keystrokes, and I’m lazy. It was well suited for a quick Python hack, and I now use the script below as a TextMate bundle so that I can just type in a list and have it converted to “the right way”.

Here is the snippet:

#!/usr/bin/env python
# encoding: utf-8


import os
import re
from sys import stdout, stdin, exit

display_list = stdin.read()
explicit = 1


def sanitize_enum(s):
    # @@ plenty to improve here
    s = s.upper().strip(' ()')
    s = re.sub('^[-\d ]+','',s)
    s = re.sub('[./() \-]', "_",s)
    while re.search('__',s):
        s = s.replace('__','_')
    return s

def format_choice(t):
    return "    ( %s, '%s' )," % t


choices = []
i = 1

for choice in display_list.split('\n'):
    if choice:
        if explicit:
            sanitized = sanitize_enum(choice)
            print "%s = %s" % (sanitized,i)
            choices.append((sanitized,choice.strip()))
        else:
            choices.append((i,choice.strip()))
        i+=1

print "choices = ("
for c in choices:
    print format_choice(c)
print "    )"