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)