from django import template
import re
import functools
register = template.Library()

# register any python function as a tag with an optional name
r_func = re.compile(r'(.*?)\s+as\s+(\w+)', re.DOTALL)
def register_function_tag(func, name=None):
    class FunctionNode(template.Node):
        def __init__(self, arg_string, var_name):
            self.arg_string = arg_string
            self.var_name = var_name

        def render(self, context):
            clist = list(context)
            clist.reverse()
            d = {"func": func}
            for c in clist:
                d.update(c)
            var = eval("func(%s)" % self.arg_string, d)
            if self.var_name is not None:
                context[self.var_name] = var 
                return ''
            else:
                return str(var)

    def do_function(parser, token):
        try:
            tag_name, arg_string = token.contents.split(None, 1)
        except ValueError:
            raise template.TemplateSyntaxError, "%r tag requires a single argument" % token.contents.split()[0]
        m = r_func.search(arg_string)
        if m is not None:
            func_string, var_name = m.groups()
        else:
            if not arg_string:
                raise template.TemplateSyntaxError, "%r tag requires one argument" % tag_name
            func_string, var_name = arg_string, None
        return FunctionNode(func_string, var_name)

    if name is not None:
        register.tag(name, do_function)
    else:
        register.tag(func.__name__, do_function)

def function_tag(name=None):
    return functools.partial(register_function_tag, name=name)


# for example
@function_tag()
def foo(arg):
    return do_something(arg)

# in template
# {% foo arg %} or
# {% foo arg as variable %} {{ variable.bar }}