def get_func_doc_for_class(name): args = [ AssignDoc(IdDoc("x"), rhs=None, annotation=IdDoc("int")), AssignDoc(IdDoc("y"), rhs=LiteralDoc(1), annotation=IdDoc("int")), ] body = [ AssignDoc(IdDoc("y"), OperationDoc(OperationKind.Add, [IdDoc("x"), LiteralDoc(1)])), AssignDoc(IdDoc("y"), OperationDoc(OperationKind.Sub, [IdDoc("y"), LiteralDoc(1)])), ] return FunctionDoc( name=IdDoc(name), args=args, decorators=[IdDoc("wrap")], return_type=LiteralDoc(None), body=body, )
def test_underline_basic(): doc = StmtBlockDoc([ ExprStmtDoc(make_id_doc("foo")), ExprStmtDoc( OperationDoc( OperationKind.Add, [make_id_doc("bar"), make_id_doc("baz")])), ExprStmtDoc(make_id_doc("qux")), ]) assert to_python_script( doc, path_to_underline=make_path("baz")) == format_script(""" foo bar + baz ^^^ qux """)
def test_underline_multiple_spans(): doc = StmtBlockDoc([ ExprStmtDoc(make_id_doc("foo")), ExprStmtDoc(make_id_doc("bar")), ExprStmtDoc( OperationDoc( OperationKind.Add, [make_id_doc("foo"), make_id_doc("foo")])), ]) assert to_python_script( doc, path_to_underline=make_path("foo")) == format_script(""" foo ^^^ bar foo + foo ^^^ ^^^ """)
def if_then_else(a, b, c): return OperationDoc(OperationKind.IfThenElse, [a, b, c])
def test_print_binary_operation_doc(op_kind, expected_token): doc = OperationDoc(op_kind, [IdDoc("x"), IdDoc("y")]) assert to_python_script(doc) == format_script(f"x {expected_token} y")
def eq(a, b): return OperationDoc(OperationKind.Eq, [a, b])
def not_eq(a, b): return OperationDoc(OperationKind.NotEq, [a, b])
def bit_xor(a, b): return OperationDoc(OperationKind.BitXor, [a, b])
def lt(a, b): return OperationDoc(OperationKind.Lt, [a, b])
def invert(a): return OperationDoc(OperationKind.Invert, [a])
def sub(a, b): return OperationDoc(OperationKind.Sub, [a, b])
def pow(a, b): return OperationDoc(OperationKind.Pow, [a, b])
def lshift(a, b): return OperationDoc(OperationKind.LShift, [a, b])
def mod(a, b): return OperationDoc(OperationKind.Mod, [a, b])
def div(a, b): return OperationDoc(OperationKind.Div, [a, b])
def mult(a, b): return OperationDoc(OperationKind.Mult, [a, b])
def test_print_special_operation_doc(op_kind, operands, expected): doc = OperationDoc(op_kind, operands) assert to_python_script(doc) == format_script(expected)
def negative(a): return OperationDoc(OperationKind.USub, [a])
def bit_and(a, b): return OperationDoc(OperationKind.BitAnd, [a, b])
def add(a, b): return OperationDoc(OperationKind.Add, [a, b])
def bit_or(a, b): return OperationDoc(OperationKind.BitOr, [a, b])
[], """ @wrap def func(x: int, y: int = 1) -> None: pass """, ), ( [ AssignDoc(IdDoc("x"), rhs=None, annotation=IdDoc("int")), AssignDoc(IdDoc("y"), rhs=LiteralDoc(1), annotation=IdDoc("int")), ], [IdDoc("wrap")], LiteralDoc(None), [ AssignDoc(IdDoc("y"), OperationDoc(OperationKind.Add, [IdDoc("x"), LiteralDoc(1)])), AssignDoc(IdDoc("y"), OperationDoc(OperationKind.Sub, [IdDoc("y"), LiteralDoc(1)])), ], """ @wrap def func(x: int, y: int = 1) -> None: y = x + 1 y = y - 1 """, ), ], ids=itertools.count(), ) def test_print_function_doc(args, decorators, body, return_type, expected): doc = FunctionDoc(IdDoc("func"), args, decorators, return_type, body) assert to_python_script(doc) == format_script(expected) # test