{% tree %}: деревья в Джанго-шаблонах

Буквально вчера ночью, вместо того, чтобы рисовать слайды для выступления на 404fest я как-то неожиданно написал кусок кода, который может кому-то ещё пригодиться. Это шаблонный тег и вспомогательный фильтр, которые умеют выводить в шаблон древовидные структуры в виде вложенных списков <ul> (вот таких, например).

Несмотря на то, что изначально я написал его "неожиданно", потом долго сидел и боролся с не очень хотящей думать головой, чтобы на этот код не противно было смотреть. И я победил! Так что код теперь production-качества, что в переводе означает "уже один раз исправлялся".

Тег {% tree %}

На самом деле, в Джанге уже есть встроенное средство для вывода древовидных структур в список: фильтр "unordered_list". Но он не подошёл, потому что умеет выводить только строчки, а мне хочется в шаблоне как-то влиять на то, что лежит в отдельных <li>: ссылку сделать, дописать что-то. Поэтому я сделал тег, который работает так:

{% tree some_list %}
<a href="{{ item.get_absolute_url }}">{{ item }}</a>
{% endtree %}

• Переменная "some_list" — структурированное дерево в том же формате, что принимает и "unordererd_list" (кому лень было ходить по ссылке, это вот такое: ['item 1', ['subitem 1', 'subitem 2'], 'item 2', 'item 3]).

• Между открывающим и закрывающим тегами находится микрошаблон, который используется для одиночного элемента списка. В него передаются две переменные:

  • "item": элемент списка
  • "level": уровень вложенности (0, 1, 2 ...)

Работающий код тега с комментариями, как говорится, for your reading pleasure:

class TreeNode(template.Node):
    def __init__(self, tree, node_list):
        self.tree = tree
        self.node_list = node_list

    def render(self, context):
        tree = self.tree.resolve(context)

        # итератор по входному списку, выдающий пары вида 
        # (элемент списка, его подсписок), причём одного из элемента пары
        # может не быть
        def pairs(items):

            # внутренний "грязный" генератор, выдающий пары, где могут быть
            # бесполезные: с обоими пустыми head и tail
            def dirty(items):
                items = iter(items)
                head = None
                try:
                    while True:
                        item = items.next()
                        if isinstance(item, (list, tuple)):
                            yield head, item
                            head = None
                        else:
                            yield head, None
                            head = item
                except StopIteration:
                    yield head, None

            # фильтр над грязным генератором, удаляющий бесполезные пары
            return ((h, t) for h, t in dirty(items) if h or t)

        # выводит элемент списка с подсписком
        # для подсписка рекурсивно вызывается render_items
        def render_item(item, sub_items, level):
            return ''.join([
                '<li>',
                item and self.node_list.render(template.Context({'item': item, 'level': level})) or '',
                sub_items and '<ul>%s</ul>' % ''.join(render_items(sub_items, level + 1)) or '',
                '</li>'
            ])

        # вывод списка элементов
        def render_items(items, level):
            return ''.join(render_item(h, t, level) for h, t in pairs(items))

        return render_items(tree, 0)

@register.tag
def tree(parser, token):
    bits = token.split_contents()
    if len(bits) != 2:
        raise template.TemplateSyntaxError('"%s" takes one argument: tree-structured list' % bits[0])
    node_list = parser.parse('end' + bits[0])
    parser.delete_first_token()
    return TreeNode(parser.compile_filter(bits[1]), node_list)

Сначала попытался написать по-тупому, разбирая входной список и подпихивая между делом всякие <li>, </li>, <ul> и прочее. Но довольно быстро стало очевидно, что такой подход приводит к невоспринимаемому вороху special-case'ов. Потому что по сути тут нужен итератор с памятью на один элемент: когда вытаскивается одиночный элемент, его нельзя тут же рендерить и закрывать, потому что за ним может идти его подсписок. А может не идти. А может итератор закончится (это особенно гадко). Поэтому я в итоге родил вот тот самый генератор "pairs", отвязанный от рендеринга. На чём и успокоился.

И да, мне это нравится больше, чем код unordered_list :-).

Фильтр astree

Ещё одна небольшая проблема — это составление списка, который принимает {% tree %}. Изначально такого списка у меня нет, а есть моделька, которая ссылается на себя же в поле "parent", и таким образом определяет дерево. Соответственно, я написал ещё и фильтр, который из queryset'а этой модели делает такой сгруппированный список:

@register.filter
def astree(items, attribute):

    # перевод списка в dict: parent -> список детей
    parent_map = defaultdict(list)
    for item in items:
        parent_map[getattr(item, attribute)].append(item)

    # рекурсивный вывод детей одного parent'а
    def tree_level(parent):
        for item in parent_map[parent]:
            yield item
            sub_items = list(tree_level(item))
            if sub_items:
                yield sub_items

    return list(tree_level(None))

Соответственно, в шаблон передаётся обычный queryset, и всё выглядит так:

{% tree object_list|astree:"parent" %}...{% endtree %}