django-mptt enabled FilteredSelectMultiple m2m widget

If you are using django-mptt to manage content (eg heirarchical categories) then it needs a bit of help to make a nice admin interface. For many-to-many fields, Django provides the quite nice FilteredSelectMultiple widget (a two-pane selection list with search box) but it only renders 'flat' lists... if you have a big category tree it's going to be confusing to know what belongs to what. Also, list items are sorted alphabetically in the js, which won't be what you want.

This snippet extends FilteredSelectMultiple to show the tree structure in the list.

fromitertoolsimportchainfromdjangoimportformsfromdjango.confimportsettingsfromdjango.contrib.adminimportwidgetsfromdjango.utils.encodingimportsmart_unicode,force_unicodefromdjango.utils.safestringimportmark_safefromdjango.utils.htmlimportescape,conditional_escapeclassMPTTModelChoiceIterator(forms.models.ModelChoiceIterator):defchoice(self,obj):tree_id=getattr(obj,getattr(self.queryset.model._meta,'tree_id_atrr','tree_id'),0)left=getattr(obj,getattr(self.queryset.model._meta,'left_atrr','lft'),0)returnsuper(MPTTModelChoiceIterator,self).choice(obj)+((tree_id,left),)classMPTTModelMultipleChoiceField(forms.ModelMultipleChoiceField):deflabel_from_instance(self,obj):level=getattr(obj,getattr(self.queryset.model._meta,'level_attr','level'),0)returnu'%s%s'%('-'*level,smart_unicode(obj))def_get_choices(self):ifhasattr(self,'_choices'):returnself._choicesreturnMPTTModelChoiceIterator(self)choices=property(_get_choices,forms.ChoiceField._set_choices)classMPTTFilteredSelectMultiple(widgets.FilteredSelectMultiple):def__init__(self,verbose_name,is_stacked,attrs=None,choices=()):super(MPTTFilteredSelectMultiple,self).__init__(verbose_name,is_stacked,attrs,choices)defrender_options(self,choices,selected_choices):""" this is copy'n'pasted from django.forms.widgets Select(Widget) change to the for loop and render_option so they will unpack and use our extra tuple of mptt sort fields (if you pass in some default choices for this field, make sure they have the extra tuple too!) """defrender_option(option_value,option_label,sort_fields):option_value=force_unicode(option_value)selected_html=(option_valueinselected_choices)andu' selected="selected"'or''returnu'<option value="%s" data-tree-id="%s" data-left-value="%s"%s>%s</option>'%(escape(option_value),sort_fields[0],sort_fields[1],selected_html,conditional_escape(force_unicode(option_label)),)# Normalize to strings.selected_choices=set([force_unicode(v)forvinselected_choices])output=[]foroption_value,option_label,sort_fieldsinchain(self.choices,choices):ifisinstance(option_label,(list,tuple)):output.append(u'<optgroup label="%s">'%escape(force_unicode(option_value)))foroptioninoption_label:output.append(render_option(*option))output.append(u'</optgroup>')else:output.append(render_option(option_value,option_label,sort_fields))returnu'\n'.join(output)classMedia:extend=Falsejs=(settings.ADMIN_MEDIA_PREFIX+"js/core.js",settings.MEDIA_URL+"js/mptt_m2m_selectbox.js",settings.ADMIN_MEDIA_PREFIX+"js/SelectFilter2.js",)