Steve Steve - 6 months ago 22
Python Question

How do you parse and inject additional nodes in a Jinja extension?

I am attempting to adapt the Jinja2 WithExtension to produce a generic extension for wrapping a block (followed by some more complex ones).

My objective is to support the following in templates:

{% wrap template='wrapper.html.j2' ... %}
<img src="{{ url('image:thumbnail' ... }}">
{% endwrap %}


And for wrapper.html.j2 to look like something like:

<div>
some ifs and stuff
{{ content }}
more ifs and stuff
</div>


I believe my example is most of the way there, WithExtension appears to parse the block and then append the AST representation of some
{% assign .. %}
nodes into the context of the nodes it is parsing.

So I figured I want the same thing, those assignments, followed by an include block, which I'd expect to be able to access those variables when the AST is parsed, and to pass through the block that was wrapped as a variable
content
.

I have the following thus far:

class WrapExtension(Extension):
tags = set(['wrap'])

def parse(self, parser):
node = nodes.Scope(lineno=next(parser.stream).lineno)
assignments = []
while parser.stream.current.type != 'block_end':
lineno = parser.stream.current.lineno
if assignments:
parser.stream.expect('comma')
target = parser.parse_assign_target()
parser.stream.expect('assign')
expr = parser.parse_expression()
assignments.append(nodes.Assign(target, expr, lineno=lineno))
content = parser.parse_statements(('name:endwrap',), drop_needle=True)
assignments.append(nodes.Name('content', content))
assignments.append(nodes.Include(nodes.Template('wrapper.html.j2'), True, False))
node.body = assignments
return node


However, it falls over at my
nodes.Include
line, I simply get
assert frame is None, 'no root frame allowed'
. I believe I need to pass AST to
nodes.Template
rather than a template name, but I don't really know how to parse in additional nodes for the objective of getting AST rather than string output (i.e. renderings) – nor whether this is the right approach. Am I on the right lines, any ideas on how I should go about this?

Answer

templatetags/wrap.py

class WrapExtension(jinja2.ext.Extension):
    tags = set(['wrap'])
    template = None

    def parse(self, parser):
        tag = parser.stream.current.value
        lineno = parser.stream.next().lineno
        args, kwargs = self.parse_args(parser)
        body = parser.parse_statements(['name:end{}'.format(tag)], drop_needle=True)

        return nodes.CallBlock(self.call_method('wrap', args, kwargs), [], [], body).set_lineno(lineno)

    def parse_args(self, parser):
        args = []
        kwargs = []
        require_comma = False

        while parser.stream.current.type != 'block_end':
            if require_comma:
                parser.stream.expect('comma')

            if parser.stream.current.type == 'name' and parser.stream.look().type == 'assign':
                key = parser.stream.current.value
                parser.stream.skip(2)
                value = parser.parse_expression()
                kwargs.append(nodes.Keyword(key, value, lineno=value.lineno))
            else:
                if kwargs:
                    parser.fail('Invalid argument syntax for WrapExtension tag',
                                parser.stream.current.lineno)
                args.append(parser.parse_expression())

            require_comma = True

        return args, kwargs

    @jinja2.contextfunction
    def wrap(self, context, caller, template=None, *args, **kwargs):
        return self.environment.get_template(template or self.template).render(dict(context, content=caller(), **kwargs))

base.html.j2

<h1>dsd</h1>
{% wrap template='wrapper.html.j2' %}
    {% for i in range(3) %}
        im wrapped content {{ i }}<br>
    {% endfor %}
{% endwrap %}

wrapper.html.j2

Hello im wrapper
<br>
<hr>
{{ content|safe }}
<hr>         

args/kwargs parsing get from here https://github.com/Suor/django-cacheops/blob/master/cacheops/jinja2.py


Additionally, the above can be extended to support additional tags with a default template specified as the wrapper:

templatetags/example.py

class ExampleExtension(WrapExtension):
    tags = set(['example'])
    template = 'example.html.j2'

base.html.j2

{% example otherstuff=True, somethingelse=False %}
    {% for i in range(3) %}
        im wrapped content {{ i }}<br>
    {% endfor %}
{% endexample %}