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

Question

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 1

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 %}