def test_equal_nodes():
assert equal_nodes(None, None)
x = ref('x')
y = ref('y')
z = ref('z')
ix = ampl.Indexing(x)
iy = ampl.Indexing(y)
check_equal_nodes(x, y)
check_equal_nodes(ampl.ParenExpr(x), ampl.UnaryExpr('-', x))
check_equal_nodes(ampl.UnaryExpr('-', x), ampl.UnaryExpr('+', x), ampl.UnaryExpr('-', y))
check_equal_nodes(ampl.BinaryExpr('*', x, y), ampl.BinaryExpr('+', x, y),
ampl.BinaryExpr('*', y, y), ampl.BinaryExpr('*', x, x))
check_equal_nodes(ampl.IfExpr(x, y, z), ampl.IfExpr(z, y, z),
ampl.IfExpr(x, x, z), ampl.IfExpr(x, y, y))
check_equal_nodes(ampl.CallExpr('sin', [x]), ampl.CallExpr('cos', [x]),
ampl.CallExpr('sin', [y]), ampl.CallExpr('sin', [x, y]))
check_equal_nodes(ampl.SumExpr(ix, x), ampl.SumExpr(iy, x), ampl.SumExpr(ix, y))
check_equal_nodes(ix, iy)
check_equal_nodes(ampl.Decl('var', 'x'), ampl.Decl('var', 'y'))
check_equal_nodes(ampl.IncludeStmt('data'), ampl.IncludeStmt('model'))
check_equal_nodes(ampl.DataStmt('param', 'S', ['a', 'b'], ['1', '2', '3']),
ampl.DataStmt('var', 'S', ['a', 'b'], ['1', '2', '3']),
ampl.DataStmt('param', 'T', ['a', 'b'], ['1', '2', '3']),
ampl.DataStmt('param', 'S', ['a'], ['1', '2', '3']),
ampl.DataStmt('param', 'S', ['a', 'c'], ['1', '2', '3']),
ampl.DataStmt('param', 'S', ['a', 'v'], ['1', '2']),
ampl.DataStmt('param', 'S', ['a', 'v'], ['1', '2', '4']))
assert not equal_nodes(ampl.Decl('var', 'x'), ampl.CompoundStmt([]))
def test_pretty_print():
check_print('a', ref('a'))
check_print('a[b]', ampl.SubscriptExpr('a', ref('b')))
check_print('(a)', ampl.ParenExpr(ref('a')))
check_print('-a', ampl.UnaryExpr('-', ref('a')))
check_print('a + b', ampl.BinaryExpr('+', ref('a'), ref('b')))
check_print('if a then b else c',
ampl.IfExpr(ref('a'), ref('b'), ref('c')))
check_print('if a then b',
ampl.IfExpr(ref('a'), ref('b'), None))
check_print('f(a, b)', ampl.CallExpr('f', [ref('a'), ref('b')]))
check_print('sum{s in S} x[s]', ampl.SumExpr(ampl.Indexing(ref('S'), 's'),
ampl.SubscriptExpr('x', ref('s'))))
check_print('{s in S}', ampl.Indexing(ref('S'), 's'))
check_print('{S}', ampl.Indexing(ref('S')))
check_print('= a', ampl.InitAttr(ref('a')))
check_print('in [a, b]', ampl.InAttr(ref('a'), ref('b')))
check_print('var x{S} = a;\n', ampl.Decl('var', 'x', ampl.Indexing(ref('S')),
[ampl.InitAttr(ref('a'))]))
decl = ampl.Decl('minimize', 'o')
decl.body = ampl.UnaryExpr('-', ref('x'))
check_print('minimize o: -x;\n', decl)
check_print('model;\n', ampl.IncludeStmt('model'))
param_names = ['a', 'b']
values = [str(n) for n in range(6)]
check_print(
'param:\n' +
'S:a b :=\n' +
'0 1 2\n' +
'3 4 5\n' +
';\n', ampl.DataStmt('param', 'S', param_names, values))
check_print('model;\nvar x;\n',
ampl.CompoundStmt([ampl.IncludeStmt('model'), ampl.Decl('var', 'x')]))
def test_parse_data():
check_parse('data;', ampl.IncludeStmt('data'))
check_parse(
'''
data;
param:
S: a b :=
1 2 3;
''',
ampl.IncludeStmt('data'), ampl.DataStmt('param', 'S', ['a', 'b'], ['1', '2', '3']))
def test_compound():
nodes = [ampl.IncludeStmt('model'), ampl.Decl('var', 'x')]
stmt = ampl.CompoundStmt(nodes)
assert type(stmt) == ampl.CompoundStmt
assert stmt.nodes == nodes
check_accept(stmt, 'visit_compound')
def test_include():
stmt = ampl.IncludeStmt('model')
assert type(stmt) == ampl.IncludeStmt
assert stmt.kind == 'model'
check_accept(stmt, 'visit_include')