def test_lines_with_nested_body(self, level: int): x = py.Line.INDENT_UNIT f = py.Function( "func", params=[], statements=[ py.Assignment("a", py.Literal(2)), py.IfElse( [(py.Literal(True), [py.Assignment("b", py.Literal(3))])], [ py.Assignment("b", py.Literal(4)), py.Assignment("c", py.Literal(1)), ], ), ], ) assert [str(l) for l in f.lines(level)] == [ x * level + "def func():", x * (level + 1) + "a = 2", x * (level + 1) + "if True:", x * (level + 2) + "b = 3", x * (level + 1) + "else:", x * (level + 2) + "b = 4", x * (level + 2) + "c = 1", ]
def test_lines_with_no_params_and_no_body(self, level: int): f = py.Function("f", params=[], statements=[]) x = py.Line.INDENT_UNIT assert [str(l) for l in f.lines(level)] == [ x * level + "def f():", x * (level + 1) + "pass", ]
def test_lines_with_hidden_comments(self, level: int): stmt = py.OpaqueBlock("foo", comments=["x"]) f = py.Function("f", params=[], statements=[stmt], comments=["1", "2"]) x = py.Line.INDENT_UNIT assert [str(l) for l in f.lines(level, comments=False)] == [ x * level + "def f():", *[str(l) for l in stmt.lines(level + 1, comments=False)], ]
def test_lines_with_complex_params(self): f = py.Function("f", params=["x: int", "abc: bool = True", "*z: str"], statements=[]) assert [str(l) for l in f.lines()] == [ "def f(x: int, abc: bool = True, *z: str):", py.Line.INDENT_UNIT + "pass", ]
def test_lines_with_hidden_comments(self, level: int): f = py.Function("f", params=[], statements=[], comments=["1", "2"]) d = py.Decoration("task", f, comments=["x", "y"]) x = py.Line.INDENT_UNIT assert [str(l) for l in d.lines(level, comments=False)] == [ x * level + "@task", x * level + "def f():", x * (level + 1) + "pass", ]
def test_with_a_function(self, level: int): f = py.Function("f", params=[], statements=[py.Assignment("a", py.Symbol("f"))]) d = py.Decoration("task(2)", f) x = py.Line.INDENT_UNIT assert [str(l) for l in d.lines(level)] == [ x * level + "@task(2)", *[str(l) for l in f.lines(level)], ]
def test_repr(self): stmts = [py.OpaqueBlock("raise")] assert ( repr( py.Function(name="f", params=["a"], statements=stmts, comments=["hi"])) == f"Function(name='f', params=['a'], statements={stmts!r}, comments=['hi'])" )
def test_nested_decorators(self, level: int): f = py.Function("f", params=[], statements=[py.Assignment("a", py.Symbol("f"))]) first = py.Decoration("task(2)", f) second = py.Decoration("task_seq(1)", first) x = py.Line.INDENT_UNIT assert [str(l) for l in second.lines(level)] == [ x * level + "@task_seq(1)", x * level + "@task(2)", *[str(l) for l in f.lines(level)], ]
def _locust_task(task: Union[Task, Task2]) -> py.Function: """ Transforms a Task into the Python code expected by Locust. This function is private because it does not return a complete Locust task (the @task decorator is missing) and should therefore not be used for that purpose by unsuspecting users. """ if isinstance(task, Task): # TODO: remove when Task2 has replaced Task. # See https://github.com/zalando-incubator/Transformer/issues/11. task = Task2.from_task(task) return py.Function(name=task.name, params=["self"], statements=task.statements)
def test_lines_with_simple_body(self, level: int): f = py.Function( "f", params=[], statements=[ py.OpaqueBlock("print('Hello!')"), py.OpaqueBlock("return") ], ) x = py.Line.INDENT_UNIT assert [str(l) for l in f.lines(level)] == [ x * level + "def f():", x * (level + 1) + "print('Hello!')", x * (level + 1) + "return", ]
def locust_imports() -> py.Program: is_post_1 = py.BinaryOp(py.Symbol("LOCUST_MAJOR_VERSION"), ">=", py.Literal(1),) imports_pre_1 = [ py.Import( ["HttpLocust", "TaskSequence", "TaskSet", "seq_task", "task"], source="locust", ) ] imports_post_1 = [ py.Import( ["HttpUser", "SequentialTaskSet", "TaskSet", "task"], source="locust", ), py.Assignment("HttpLocust", py.Symbol("HttpUser")), py.Assignment("TaskSequence", py.Symbol("SequentialTaskSet")), py.Function("seq_task", ["_"], [py.Return(py.Symbol("task"))]), ] return [ py.IfElse([(is_post_1, imports_post_1)], imports_pre_1), ]
def test_repr(self): f = py.Function("f", params=[], statements=[]) assert (repr(py.Decoration( "task", f, comments=["hi"])) == f"Decoration('task', {f!r}, comments=['hi'])")
def test_lines_with_simple_params(self): f = py.Function("f", params=["x", "y"], statements=[]) assert [str(l) for l in f.lines()] == [ "def f(x, y):", py.Line.INDENT_UNIT + "pass", ]