diff --git a/synr/ast.py b/synr/ast.py index 89cf2db..2c99564 100644 --- a/synr/ast.py +++ b/synr/ast.py @@ -642,6 +642,24 @@ class For(Stmt): body: Block +@attr.s(auto_attribs=True, frozen=True) +class While(Stmt): + """An while statement. + + Examples + -------- + .. code-block:: python + + while x <= 2: + pass + + Here :code:`condition` is :code:`x <= 2`, and :code:`body` will be :code:`pass`. + """ + + condition: Expr + body: Block + + @attr.s(auto_attribs=True, frozen=True) class With(Stmt): """A with statement. diff --git a/synr/compiler.py b/synr/compiler.py index 9e39a29..4db39f8 100644 --- a/synr/compiler.py +++ b/synr/compiler.py @@ -288,6 +288,11 @@ def compile_stmt(self, stmt: py_ast.stmt) -> Stmt: body = self.compile_block(stmt.body) return For(self.span_from_ast(stmt), lhs_vars, rhs, body) + elif isinstance(stmt, py_ast.While): + condition = self.compile_expr(stmt.test) + body = self.compile_block(stmt.body) + return While(stmt_span, condition, body) + elif isinstance(stmt, py_ast.With): if len(stmt.items) != 1: self.error( diff --git a/tests/test_synr.py b/tests/test_synr.py index 5b14e5d..1936a20 100644 --- a/tests/test_synr.py +++ b/tests/test_synr.py @@ -92,6 +92,26 @@ def test_for(): assert fr.body.stmts[0].value.id.name == "x" +def func_while(): + while x < 10: + return x + + +def test_while(): + module = to_ast(func_while) + fn = assert_one_fn(module, "func_while", no_params=0) + + while_stmt = fn.body.stmts[0] + assert isinstance(while_stmt, synr.ast.While) + assert isinstance(while_stmt.body.stmts[0], synr.ast.Return) + assert while_stmt.body.stmts[0].value.id.name == "x" + cond = while_stmt.condition + assert isinstance(cond, synr.ast.Call) + assert cond.func_name.name == synr.ast.BuiltinOp.LT + assert cond.params[0].id.name == "x" + assert cond.params[1].value == 10 + + def func_with(): with x as y: return x @@ -615,6 +635,7 @@ def foo(): test_id_function() test_class() test_for() + test_while() test_with() test_block() test_assign()