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