WithTag Tag

# -*- coding: utf-8 -*-
from django.template import Library, Node, TemplateSyntaxError, Token, TOKEN_BLOCK

register = Library()

class WithTagNode(Node):
    """Set a context variable with the returned value by any templatetag.

    For example::
        {% withtag url projects.views.profile user.username as user_profile_url %}
        ...
        {{ user_profile_url }}
        ...
        {% endwithtag %}

    """
    def __init__(self, tag, var_name, parser, token, nodelist):
        self.tag = tag
        self.var_name = var_name
        self.parser = parser
        self.token = token
        self.nodelist = nodelist

    def __repr__(self):
        return "<WithTagNode>"

    def render(self, context):
        templatetag = self.parser.tags.get(self.tag, None)
        if not templatetag or not callable(templatetag):
            raise TemplateSyntaxError, "Invalid block tag '%r' as argument " \
                                       "for withtag block." % self.tag
        node = templatetag(self.parser, self.token)
        context.push()
        context[self.var_name] = node.render(context)
        output = self.nodelist.render(context)
        context.pop()
        return output

def do_withtag(parser, token):
    bits = token.split_contents()
    bits_length = len(bits)
    nodelist = parser.parse(('endwithtag',))
    parser.delete_first_token()
    if bits_length > 3 and bits[bits_length-2] == 'as':
        new_token = Token(TOKEN_BLOCK, ' '.join(bits[1:bits_length-2]))
        return WithTagNode(bits[1], bits[bits_length-1], parser, new_token, nodelist)
    else:
        raise TemplateSyntaxError, "%r tag requires one \"as\" argument." \
                                   % token.contents.split()[0]

register.tag('withtag', do_withtag)