def test_special_function():
output = printer.RawStream(special_functions=True)
assert output.get_printer_for_function('foo.test_function') is None
try:
@printer.special_function('foo.test_function')
def test(node, output):
pass
assert output.get_printer_for_function('foo.test_function') is test
with pytest.raises(printer.PrinterAlreadyPresentError):
@printer.special_function('foo.test_function')
def test1(node, output):
pass
@printer.special_function('foo.test_function', override=True)
def test_function(node, output):
output.print_list(node.args, '-')
result = output('SELECT foo.test_function(x, "Y") FROM sometable')
assert result == 'SELECT x - "Y" FROM sometable'
finally:
printer.SPECIAL_FUNCTIONS.pop('foo.test_function')
def test_raw_stream_with_sql():
raw_stmt_printer = printer.NODE_PRINTERS.pop('RawStmt', None)
try:
@printer.node_printer('RawStmt')
def raw_stmt(node, output):
output.write('Yeah')
output = printer.RawStream()
result = output('SELECT 1; SELECT 2')
assert result == 'Yeah; Yeah'
finally:
if raw_stmt_printer is not None:
printer.NODE_PRINTERS['RawStmt'] = raw_stmt_printer
else:
printer.NODE_PRINTERS.pop('RawStmt', None)
def test_raw_stream_basics():
ptree = [{'TestRoot': {'bar': {'TestChild': {'a': [
{'TestNiece': {'x': 0, 'y': 0}},
{'TestNiece': {'x': 1, 'z': [
{'TestNiece': {'x': 2, 'y': 2}},
{'TestNiece': {'x': 3, 'z': [
{'TestNiece': {'x': 4, 'y': 4}},
{'TestNiece': {'x': 5, 'y': 5}},
]}}
]}},
]}}}}]
root = Node(ptree)
output = printer.RawStream()
with pytest.raises(NotImplementedError) as exc:
output(root)
assert "'TestRoot'" in str(exc)
try:
@printer.node_printer('TestRoot')
def test_root(node, output):
output.write('bar = ')
output.print_node(node.bar)
@printer.node_printer('TestChild')
def test_child(node, output):
output.write('a(')
with output.push_indent():
output.print_list(node.a, '*')
output.write(')')
@printer.node_printer('TestNiece')
def test_niece(node, output):
output.write('{')
output.write('x')
output.print_node(node.x)
output.separator()
output.write(',')
output.separator()
if node.y:
output.write('y')
output.print_node(node.y)
else:
if output.test_child_z_is_expression:
with output.push_indent(2, relative=False):
output.print_list(node.z, '+')
else:
output.print_list(node.z, '/', standalone_items=False)
output.write('}')
output = printer.RawStream()
output.test_child_z_is_expression = True
result = output(root)
assert result == 'bar = a({x0, y0} * {x1,{x2, y2} + {x3,{x4, y4} + {x5, y5}}})'
output = printer.RawStream()
output.test_child_z_is_expression = False
result = output(root)
assert result == 'bar = a({x0, y0} * {x1,{x2, y2} / {x3,{x4, y4} / {x5, y5}}})'
finally:
printer.NODE_PRINTERS.pop('TestRoot', None)
printer.NODE_PRINTERS.pop('TestChild', None)
printer.NODE_PRINTERS.pop('TestNiece', None)
def do(meth, node):
output = printer.RawStream()
getattr(output, meth)(node)
return output.getvalue()
def test_raw_stream_invalid_call():
with pytest.raises(ValueError) as exc:
printer.RawStream()(1)
