1# -*- coding: utf-8 -*-
2
3import arrow
4
5from jinja2 import nodes
6from jinja2.ext import Extension
7
8
9class TimeExtension(Extension):
10    tags = set(['now'])
11
12    def __init__(self, environment):
13        super(TimeExtension, self).__init__(environment)
14
15        # add the defaults to the environment
16        environment.extend(datetime_format='%Y-%m-%d')
17
18    def _datetime(self, timezone, operator, offset, datetime_format):
19        d = arrow.now(timezone)
20
21        # Parse replace kwargs from offset and include operator
22        replace_params = {}
23        for param in offset.split(','):
24            interval, value = param.split('=')
25            replace_params[interval.strip()] = float(operator + value.strip())
26        d = d.replace(**replace_params)
27
28        if datetime_format is None:
29            datetime_format = self.environment.datetime_format
30        return d.strftime(datetime_format)
31
32    def _now(self, timezone, datetime_format):
33        if datetime_format is None:
34            datetime_format = self.environment.datetime_format
35        return arrow.now(timezone).strftime(datetime_format)
36
37    def parse(self, parser):
38        lineno = next(parser.stream).lineno
39
40        node = parser.parse_expression()
41
42        if parser.stream.skip_if('comma'):
43            datetime_format = parser.parse_expression()
44        else:
45            datetime_format = nodes.Const(None)
46
47        if isinstance(node, nodes.Add):
48            call_method = self.call_method(
49                '_datetime',
50                [node.left, nodes.Const('+'), node.right, datetime_format],
51                lineno=lineno,
52            )
53        elif isinstance(node, nodes.Sub):
54            call_method = self.call_method(
55                '_datetime',
56                [node.left, nodes.Const('-'), node.right, datetime_format],
57                lineno=lineno,
58            )
59        else:
60            call_method = self.call_method(
61                '_now',
62                [node, datetime_format],
63                lineno=lineno,
64            )
65        return nodes.Output([call_method], lineno=lineno)
66