1import jinja2
2import jinja2.nodes
3
4from jinja2.ext import Extension
5
6from .templatetags.wagtailcore_tags import pageurl, richtext, slugurl, wagtail_version
7
8
9class WagtailCoreExtension(Extension):
10    tags = {'include_block'}
11
12    def __init__(self, environment):
13        super().__init__(environment)
14
15        self.environment.globals.update({
16            'pageurl': jinja2.contextfunction(pageurl),
17            'slugurl': jinja2.contextfunction(slugurl),
18            'wagtail_version': wagtail_version,
19        })
20        self.environment.filters.update({
21            'richtext': richtext,
22        })
23
24    def parse(self, parser):
25        parse_method = getattr(self, 'parse_' + parser.stream.current.value)
26
27        return parse_method(parser)
28
29    def parse_include_block(self, parser):
30        lineno = next(parser.stream).lineno
31
32        args = [parser.parse_expression()]
33
34        # always pass context to _include_block - even if we're not passing it on to render_as_block,
35        # we need it to see if autoescaping is enabled
36        if hasattr(jinja2.nodes, 'DerivedContextReference'):
37            # DerivedContextReference includes local variables. Introduced in Jinja 2.11
38            args.append(jinja2.nodes.DerivedContextReference())
39        else:
40            args.append(jinja2.nodes.ContextReference())
41
42        use_context = True
43        if parser.stream.current.test_any('name:with', 'name:without') and parser.stream.look().test('name:context'):
44            use_context = next(parser.stream).value == 'with'
45            parser.stream.skip()
46
47        args.append(jinja2.nodes.Const(use_context))
48
49        node = self.call_method('_include_block', args, lineno=lineno)
50        return jinja2.nodes.Output([node], lineno=lineno)
51
52    def _include_block(self, value, context, use_context):
53        if hasattr(value, 'render_as_block'):
54            if use_context:
55                new_context = context.get_all()
56            else:
57                new_context = {}
58
59            result = value.render_as_block(context=new_context)
60        else:
61            result = value
62
63        if context.eval_ctx.autoescape:
64            return jinja2.escape(result)
65        else:
66            return jinja2.Markup(result)
67
68
69# Nicer import names
70core = WagtailCoreExtension
71