Jinja2 extension for Flask + FormEncode

form.py

from flask import current_app, request, url_for
from formencode.htmlfill import render
from formencode.api import Invalid
from jinja2 import Markup, Undefined
from jinja2.exceptions import TemplateSyntaxError
from jinja2.ext import Extension
from jinja2.nodes import (Const, Filter, FilterBlock, Keyword, Output,
                          TemplateData)
from werkzeug.utils import escape

__all__ = 'FormExtension',


class FormExtension(Extension):
    """The ``form`` tag for Jinja templates.

    .. sourcecode:: html+jinja

       {% form 'comic.delete_comic', comic_id=comic.id %}
         <button type="submit">Delete</button>
       {% endform %}

    The above code will be compiled into:

    .. sourcecode:: html

       <form action="/comics/1234?_method=DELETE" method="POST">
         <button type="submit">Delete</button>
       </form>

    Note that the compiled HTML's ``form`` tag also applies
    :class:`~crosspop.web.util.MethodRewriteMiddleware` also.

    You can specify form attributes as well by passing dictionary:

    .. sourcecode:: html+jinja

       {% form 'comic.delete_comic', comic_id=comic.id,
               {'id': 'delete-comic', 'enctype': 'multipart/form-data'} %}
         <button type="submit">Delete</button>
       {% endform %}

    It also does :mod:`~formencode.htmlfill`, and you can specify
    defaults and errors using its optional ``with`` and ``else``
    keywords:

    .. sourcecode:: html+jinja

       {% form 'comic.add_comic', comic_id=comic.id with {'title': 'Title'} %}
         <label>Title <input type="text" name="title"></label>
         <button type="submit">Save</button>
       {% endform %}

    The above code will turn into:

    .. sourcecode:: html

       <form action="/comics/" method="POST">
         <label>Title <input type="text" name="title" value="Title"></label>
         <button type="submit">Save</button>
       </form>

    If you append ``else`` keyword and following errors dictionary,
    you can use ``iferror`` and ``error`` tags inside ``form`` tag:

    .. sourcecode:: html+jinja

       {% form 'comic.add_comic', comic_id=comic.id
          with {'title': 'Title'} else {'title': "It's empty!"} %}
         <label>Title
           <input type="text" name="title">
           {% iferror 'title' -%}
             <span class="error"> {%- error -%} </span>
           {%- endiferror -%}
         </label>
         <button type="submit">Save</button>
       {% endform %}

    The above code will turn into:

    .. sourcecode:: html

       <form action="/comics/" method="POST">
         <label>Title
           <input type="text" name="title" value="Title">
           <span class="error">It's empty!</span></label>
         <button type="submit">Save</button>
       </form>

    You can use only ``error`` tag without ``iferror`` tag,
    and in that case it has the name e.g.::

    .. sourcecode:: html+jinja

       {% error 'title' %}

    The above code is equivalent to:

    .. sourcecode:: html+jinja

       {% iferror 'title' -%}
         {%- error -%}
       {%- endif %}

    """

    tags = set(['form'])

    def __init__(self, environment):
        super(FormExtension, self).__init__(environment)
        environment.filters['FormExtension$htmlfill'] = self.htmlfill

    def parse(self, parser):
        """Parser hook to expand ``form`` tag."""
        try:
            return self._parse(parser)
        except:
            import traceback; traceback.print_exc()
            raise

    def _parse(self, parser):
        parser.stream.next()
        endpoint = parser.parse_expression()
        endpoint_args = []
        form_attrs = []
        if parser.stream.skip_if('comma'):
            while not (parser.stream.current.type == 'block_end' or
                       parser.stream.current.type == 'name' and
                       parser.stream.current.value == 'with'):
                if endpoint_args or form_attrs:
                    parser.stream.expect('comma')
                name = parser.stream.next_if('name')
                if name and parser.stream.skip_if('assign'):
                    keyword = Keyword(
                        name.value,
                        parser.parse_expression()
                    )
                    endpoint_args.append(keyword)
                else:
                    if name:
                        parser.stream.push(parser.stream.current)
                        parser.stream.current = name
                    form_attrs.append(parser.parse_expression())
        if parser.stream.next_if('name:with'):
            htmlfill_defaults = parser.parse_expression()
            if parser.stream.next_if('name:else'):
                htmlfill_errors = parser.parse_expression()
            else:
                htmlfill_errors = Const(None)
        else:
            htmlfill_defaults = Const(None)
            htmlfill_errors = Const(None)
        lineno = parser.stream.current.lineno
        body = []
        current_error_name = None
        while 1:
            body += parser.parse_statements([
                'name:endform', 'name:iferror',
                'name:endiferror', 'name:error'
            ])
            if parser.stream.current.type == 'name':
                tag_name = parser.stream.current.value
                if tag_name == 'endform':
                    next(parser.stream)
                    break
                elif tag_name == 'iferror':
                    next(parser.stream)
                    current_error_name = parser.parse_expression()
                    output = [
                        TemplateData('<form:iferror name="'),
                        current_error_name,
                        TemplateData('">')
                    ]
                elif tag_name == 'endiferror':
                    current_error_name = None
                    output = [TemplateData('</form:iferror>')]
                    next(parser.stream)
                elif tag_name == 'error':
                    if (current_error_name and
                        parser.stream.look().type == 'block_end'):
                        error_name = current_error_name
                        next(parser.stream)
                    else:
                        next(parser.stream)
                        error_name = parser.parse_expression()
                    output = [
                        TemplateData('<form:error name="'),
                        error_name,
                        TemplateData('" format="escape">'),
                    ]
                else:
                    raise TemplateSyntaxError(
                        'Encountered unknown tag ' + repr(tag_name)
                    )
                body.append(Output(output))
        tag = self.call_method('_form_tag', [endpoint] + form_attrs,
                               endpoint_args, lineno=lineno)
        filter_block = FilterBlock(lineno=lineno)
        filter_block.body = body
        filter_block.filter = Filter(
            None, 'FormExtension$htmlfill',
            [htmlfill_defaults, htmlfill_errors], {}, None, None,
            lineno=lineno
        )
        return [
            Output([tag]),
            filter_block,
            Output([TemplateData(u'</form>')])
        ]

    def get_method_for(self, endpoint):
        """Gets the available method for ``endpoint``.

        :param endpoint: an endpoint (view function) name to find
                         the available method
        :type endpoint: :class:`basestring`
        :returns: an available method
        :rtype: :class:`basestring`

        """
        if endpoint.startswith('.'):
            mod = request.blueprint
            if mod is None:
                endpoint = endpoint[1:]
            else:
                endpoint = mod + endpoint
        elif endpoint.startswith('.'):
            endpoint = endpoint[1:]
        methods = current_app.url_map._rules_by_endpoint[endpoint][0].methods \
                             .difference(['HEAD', 'OPTIONS'])
        for method in 'POST', 'PUT', 'DELETE', 'GET':
            if method in methods:
                return method
        return next(methods, None)

    def htmlfill(self, form, defaults, errors):
        if isinstance(defaults, Undefined):
            defaults = None
        if isinstance(errors, Undefined):
            errors = None
        if isinstance(defaults, Invalid) and errors is None:
            defaults, errors = defaults.value, defaults.error_dict
        if isinstance(errors, Invalid):
            errors = errors.error_dict
        return render(form, defaults, errors, auto_insert_errors=False)

    def _form_tag(self, endpoint, *attr_dicts, **values):
        method = self.get_method_for(endpoint)
        action = url_for(endpoint, _method=method, **values)
        if method not in ('GET', 'POST'):
            action += ('&' if '?' in action else '?') + '_method=' + method
        form_method = 'GET' if method == 'GET' else 'POST'
        attrs = {'action': action, 'method': form_method}
        for attr_dict in attr_dicts:
            attrs.update(attr_dict)
        pairs = attrs.items()
        pairs.sort(key=lambda (k, v): k)
        attrs_str = u' '.join(u'{0}="{1}"'.format(k, escape(v))
                              for k, v in pairs)
        return Markup(u'<form {0}>'.format(attrs_str))

form_test.py

from flask import Flask, render_template_string

from form import FormExtension

app = Flask(__name__)
app.jinja_env.add_extension(FormExtension)


class ValueStore(object):

    def __init__(self):
        self.stored = False
        self.value = None

    def __call__(self, value):
        self.value = value
        self.stored = True
        return 'STORED'


def test_form_extension():
    @app.route('/__test__/form_extension/<int:a>/<b>', methods=['DELETE'])
    def _test_form_extension(a, b):
        return ''
    # GET
    with app.test_request_context('/'):
        html = render_template_string('''
            {%- autoescape true -%}
                {%- form 'user.signup_form' -%}
                    body
                {%- endform -%}
            {%- endautoescape -%}
        ''')
    assert html == '<form action="/signup/" method="GET">body</form>'
    # POST
    with app.test_request_context('/'):
        html = render_template_string('''
            {%- autoescape true -%}
                {%- form 'user.signup' -%}
                    body
                {%- endform -%}
            {%- endautoescape -%}
        ''')
    assert html == '<form action="/" method="POST">body</form>'
    # PUT/DELETE
    with app.test_request_context('/'):
        html = render_template_string('''
            {%- autoescape true -%}
                {%- form '_test_form_extension', a=123, b='hi' -%}
                    body
                {%- endform -%}
            {%- endautoescape -%}
        ''')
    expected = ('<form action="/__test__/form_extension/123/hi?_method=DELETE"'
                ' method="POST">body</form>')
    assert html == expected
    # form attributes
    with app.test_request_context('/'):
        html = render_template_string('''
            {%- autoescape true -%}
                {%- form 'user.signup', {'class': 'form', 'id': 'signup'} -%}
                    body
                {%- endform -%}
            {%- endautoescape -%}
        ''')
    expected = ('<form action="/" class="form" id="signup" '
                'method="POST">body</form>')
    assert html == expected
    # htmlfill defaults
    with app.test_request_context('/'):
        html = render_template_string('''
            {%- autoescape true -%}
                {%- form 'user.signup' with {'a': 123, 'b': 456} -%}
                    <input type="number" name="a"> {{- '' -}}
                    <input type="number" name="b" value="2"> {{- '' -}}
                    <input type="number" name="c" value="3">
                {%- endform -%}
            {%- endautoescape -%}
        ''')
    expected = ('<form action="/" method="POST">'
                '<input type="number" name="a" value="123">'
                '<input type="number" name="b" value="456">'
                '<input type="number" name="c" value=""></form>')
    assert html == expected
    # htmlfill defaults + errors
    with app.test_request_context('/'):
        html = render_template_string('''
            {%- autoescape true -%}
                {%- form 'user.signup' with {'a': 123, 'b': 456}
                    else {'a': '<b>error</b> message', 'b': 'error <b>'} -%}
                    <input type="number" name="a">
                    {%- error 'a' -%}
                    <input type="number" name="b" value="2">
                    {%- iferror 'b' -%}
                        <span> {%- error -%} </span>
                    {%- endiferror -%}
                    <input type="number" name="c" value="3">
                    {%- iferror 'c' -%}
                        <span> {%- error -%} </span>
                    {%- endiferror -%}
                {%- endform -%}
            {%- endautoescape -%}
        ''')
    expected = ('<form action="/" method="POST">'
                '<input type="number" name="a" class="error" value="123">'
                '&lt;b&gt;error&lt;/b&gt; message'
                '<input type="number" name="b" value="456" class="error">'
                '<span>error &lt;b&gt;</span>'
                '<input type="number" name="c" value=""></form>')
    assert html == expected
    # form attributes + htmlfill defaults
    with app.test_request_context('/'):
        html = render_template_string('''
            {%- autoescape true -%}
                {%- form 'user.signup', {'class': 'form', 'id': 'signup'}
                    with {'a': 123, 'b': 456} -%}
                    <input type="number" name="a"> {{- '' -}}
                    <input type="number" name="b" value="2"> {{- '' -}}
                    <input type="number" name="c" value="3">
                {%- endform -%}
            {%- endautoescape -%}
        ''')
    expected = ('<form action="/" class="form" id="signup" '
                'method="POST"><input type="number" name="a" value="123">'
                '<input type="number" name="b" value="456">'
                '<input type="number" name="c" value=""></form>')
    assert html == expected