from jinja2 import nodes
from jinja2.idtracking import symbols_for_node


def test_basics():
    for_loop = nodes.For(
        nodes.Name('foo', 'store'),
        nodes.Name('seq', 'load'),
        [nodes.Output([nodes.Name('foo', 'load')])],
        [], None, False)
    tmpl = nodes.Template([
        nodes.Assign(
            nodes.Name('foo', 'store'),
            nodes.Name('bar', 'load')),
        for_loop])

    sym = symbols_for_node(tmpl)
    assert sym.refs == {
        'foo': 'l_0_foo',
        'bar': 'l_0_bar',
        'seq': 'l_0_seq',
    }
    assert sym.loads == {
        'l_0_foo': ('undefined', None),
        'l_0_bar': ('resolve', 'bar'),
        'l_0_seq': ('resolve', 'seq'),
    }

    sym = symbols_for_node(for_loop, sym)
    assert sym.refs == {
        'foo': 'l_1_foo',
    }
    assert sym.loads == {
        'l_1_foo': ('param', None),
    }


def test_complex():
    title_block = nodes.Block('title', [
        nodes.Output([nodes.TemplateData(u'Page Title')])
    ], False)

    render_title_macro = nodes.Macro('render_title', [nodes.Name('title', 'param')], [], [
        nodes.Output([
            nodes.TemplateData(u'\n  <div class="title">\n    <h1>'),
            nodes.Name('title', 'load'),
            nodes.TemplateData(u'</h1>\n    <p>'),
            nodes.Name('subtitle', 'load'),
            nodes.TemplateData(u'</p>\n    ')]),
        nodes.Assign(
            nodes.Name('subtitle', 'store'), nodes.Const('something else')),
        nodes.Output([
            nodes.TemplateData(u'\n    <p>'),
            nodes.Name('subtitle', 'load'),
            nodes.TemplateData(u'</p>\n  </div>\n'),
            nodes.If(
                nodes.Name('something', 'load'), [
                    nodes.Assign(nodes.Name('title_upper', 'store'),
                                 nodes.Filter(nodes.Name('title', 'load'),
                                              'upper', [], [], None, None)),
                    nodes.Output([
                        nodes.Name('title_upper', 'load'),
                        nodes.Call(nodes.Name('render_title', 'load'), [
                            nodes.Const('Aha')], [], None, None)])], [], [])])])

    for_loop = nodes.For(
        nodes.Name('item', 'store'),
        nodes.Name('seq', 'load'), [
            nodes.Output([
                nodes.TemplateData(u'\n    <li>'),
                nodes.Name('item', 'load'),
                nodes.TemplateData(u'</li>\n    <span>')]),
            nodes.Include(nodes.Const('helper.html'), True, False),
            nodes.Output([
                nodes.TemplateData(u'</span>\n  ')])], [], None, False)

    body_block = nodes.Block('body', [
        nodes.Output([
            nodes.TemplateData(u'\n  '),
            nodes.Call(nodes.Name('render_title', 'load'), [
                nodes.Name('item', 'load')], [], None, None),
            nodes.TemplateData(u'\n  <ul>\n  ')]),
        for_loop,
        nodes.Output([nodes.TemplateData(u'\n  </ul>\n')])],
        False)

    tmpl = nodes.Template([
        nodes.Extends(nodes.Const('layout.html')),
        title_block,
        render_title_macro,
        body_block,
    ])

    tmpl_sym = symbols_for_node(tmpl)
    assert tmpl_sym.refs == {
        'render_title': 'l_0_render_title',
    }
    assert tmpl_sym.loads == {
        'l_0_render_title': ('undefined', None),
    }
    assert tmpl_sym.stores == set(['render_title'])
    assert tmpl_sym.dump_stores() == {
        'render_title': 'l_0_render_title',
    }

    macro_sym = symbols_for_node(render_title_macro, tmpl_sym)
    assert macro_sym.refs == {
        'subtitle': 'l_1_subtitle',
        'something': 'l_1_something',
        'title': 'l_1_title',
        'title_upper': 'l_1_title_upper',
    }
    assert macro_sym.loads == {
        'l_1_subtitle': ('resolve', 'subtitle'),
        'l_1_something': ('resolve','something'),
        'l_1_title': ('param', None),
        'l_1_title_upper': ('resolve', 'title_upper'),
    }
    assert macro_sym.stores == set(['title', 'title_upper', 'subtitle'])
    assert macro_sym.find_ref('render_title') == 'l_0_render_title'
    assert macro_sym.dump_stores() == {
        'title': 'l_1_title',
        'title_upper': 'l_1_title_upper',
        'subtitle': 'l_1_subtitle',
        'render_title': 'l_0_render_title',
    }

    body_sym = symbols_for_node(body_block)
    assert body_sym.refs == {
        'item': 'l_0_item',
        'seq': 'l_0_seq',
        'render_title': 'l_0_render_title',
    }
    assert body_sym.loads == {
        'l_0_item': ('resolve', 'item'),
        'l_0_seq': ('resolve', 'seq'),
        'l_0_render_title': ('resolve', 'render_title'),
    }
    assert body_sym.stores == set([])

    for_sym = symbols_for_node(for_loop, body_sym)
    assert for_sym.refs == {
        'item': 'l_1_item',
    }
    assert for_sym.loads == {
        'l_1_item': ('param', None),
    }
    assert for_sym.stores == set(['item'])
    assert for_sym.dump_stores() == {
        'item': 'l_1_item',
    }


def test_if_branching_stores():
    tmpl = nodes.Template([
        nodes.If(nodes.Name('expression', 'load'), [
            nodes.Assign(nodes.Name('variable', 'store'),
                         nodes.Const(42))], [], [])])

    sym = symbols_for_node(tmpl)
    assert sym.refs == {
        'variable': 'l_0_variable',
        'expression': 'l_0_expression'
    }
    assert sym.stores == set(['variable'])
    assert sym.loads == {
        'l_0_variable': ('resolve', 'variable'),
        'l_0_expression': ('resolve', 'expression')
    }
    assert sym.dump_stores() == {
        'variable': 'l_0_variable',
    }


def test_if_branching_stores_undefined():
    tmpl = nodes.Template([
        nodes.Assign(nodes.Name('variable', 'store'), nodes.Const(23)),
        nodes.If(nodes.Name('expression', 'load'), [
            nodes.Assign(nodes.Name('variable', 'store'),
                         nodes.Const(42))], [], [])])

    sym = symbols_for_node(tmpl)
    assert sym.refs == {
        'variable': 'l_0_variable',
        'expression': 'l_0_expression'
    }
    assert sym.stores == set(['variable'])
    assert sym.loads == {
        'l_0_variable': ('undefined', None),
        'l_0_expression': ('resolve', 'expression')
    }
    assert sym.dump_stores() == {
        'variable': 'l_0_variable',
    }


def test_if_branching_multi_scope():
    for_loop = nodes.For(nodes.Name('item', 'store'), nodes.Name('seq', 'load'), [
        nodes.If(nodes.Name('expression', 'load'), [
            nodes.Assign(nodes.Name('x', 'store'), nodes.Const(42))], [], []),
        nodes.Include(nodes.Const('helper.html'), True, False)
    ], [], None, False)

    tmpl = nodes.Template([
        nodes.Assign(nodes.Name('x', 'store'), nodes.Const(23)),
        for_loop
    ])

    tmpl_sym = symbols_for_node(tmpl)
    for_sym = symbols_for_node(for_loop, tmpl_sym)
    assert for_sym.stores == set(['item', 'x'])
    assert for_sym.loads == {
        'l_1_x': ('alias', 'l_0_x'),
        'l_1_item': ('param', None),
        'l_1_expression': ('resolve', 'expression'),
    }
