feat: add ContinueStmt and BreakStmt (#255)

Author: Sarfia-786

libs/astx-transpilers/src/astx_transpilers/python_string.py

@@ -115,6 +115,11 @@ def visit(self, node: astx.Block) -> str:
         """Handle Block nodes."""
         return self._generate_block(node)
 
+    @dispatch  # type: ignore[no-redef]
+    def visit(self, node: astx.BreakStmt) -> str:
+        """Handle BreakStmt nodes."""
+        return "break"
+
     @dispatch  # type: ignore[no-redef]
     def visit(self, node: astx.CaseStmt) -> str:
         """Handle CaseStmt nodes."""
@@ -154,6 +159,11 @@ def visit(self, node: astx.ClassDefStmt) -> str:
         class_type = "(ABC)" if node.is_abstract else ""
         return f"class {node.name}{class_type}:\n{self.visit(node.body)}"
 
+    @dispatch  # type: ignore[no-redef]
+    def visit(self, node: astx.ContinueStmt) -> str:
+        """Handle ContinueStmt nodes."""
+        return "continue"
+
     @dispatch  # type: ignore[no-redef]
     def visit(self, node: astx.ComprehensionClause) -> str:
         """Handle ComprehensionClause nodes."""

libs/astx-transpilers/tests/test_python.py

@@ -457,6 +457,52 @@ def test_transpiler_async_for_range_loop_expr() -> None:
     )
 
 
+def test_transpiler_break_stmt() -> None:
+    """Test astx.BreakStmt transpilation."""
+    # Create a simple loop structure (e.g., WhileStmt)
+    x_var = astx.Variable(name="x")
+    condition = astx.BinaryOp(op_code="<", lhs=x_var, rhs=astx.LiteralInt32(5))
+
+    # Create the loop body with a break statement
+    body_block = astx.Block(name="while_body")
+    body_block.append(astx.BreakStmt())
+
+    while_stmt = astx.WhileStmt(condition=condition, body=body_block)
+
+    # Generate Python code
+    generated_code = translate(while_stmt)
+
+    # Expected code for the WhileStmt with break
+    expected_code = "while (x < 5):\n    break"
+
+    assert generated_code == expected_code, (
+        f"Expected '{expected_code}', but got '{generated_code}'"
+    )
+
+
+def test_transpiler_continue_stmt() -> None:
+    """Test astx.ContinueStmt transpilation."""
+    # Create a simple loop structure (e.g., WhileStmt)
+    x_var = astx.Variable(name="x")
+    condition = astx.BinaryOp(op_code="<", lhs=x_var, rhs=astx.LiteralInt32(5))
+
+    # Create the loop body with a continue statement
+    body_block = astx.Block(name="while_body")
+    body_block.append(astx.ContinueStmt())
+
+    while_stmt = astx.WhileStmt(condition=condition, body=body_block)
+
+    # Generate Python code
+    generated_code = translate(while_stmt)
+
+    # Expected code for the WhileStmt with continue
+    expected_code = "while (x < 5):\n    continue"
+
+    assert generated_code == expected_code, (
+        f"Expected '{expected_code}', but got '{generated_code}'"
+    )
+
+
 def test_transpiler_binary_op() -> None:
     """Test astx.BinaryOp for addition operation."""
     # Create a BinaryOp node for the expression "x + y"

libs/astx/src/astx/__init__.py

@@ -69,7 +69,9 @@
 from astx.flows import (
     AsyncForRangeLoopExpr,
     AsyncForRangeLoopStmt,
+    BreakStmt,
     CaseStmt,
+    ContinueStmt,
     DoWhileExpr,
     DoWhileStmt,
     ForCountLoopExpr,
@@ -224,6 +226,7 @@ def get_version() -> str:
     "BoolBinaryOp",
     "BoolUnaryOp",
     "Boolean",
+    "BreakStmt",
     "CaseStmt",
     "CatchHandlerStmt",
     "ClassDeclStmt",
@@ -235,6 +238,7 @@ def get_version() -> str:
     "Complex64",
     "Comprehension",
     "ComprehensionClause",
+    "ContinueStmt",
     "DataType",
     "DataTypeOps",
     "Date",

libs/astx/src/astx/base.py

@@ -140,6 +140,8 @@ class ASTKind(Enum):
     DoWhileStmtKind = -515
     DoWhileExprKind = -516
     GeneratorExprKind = -517
+    BreakStmtKind = -518
+    ContinueStmtKind = -519
 
     # data types
     NullDTKind = -600

libs/astx/src/astx/flows.py

@@ -472,6 +472,56 @@ def get_struct(self, simplified: bool = False) -> ReprStruct:
         return self._prepare_struct(key, value, simplified)
 
 
+@public
+@typechecked
+class BreakStmt(StatementType):
+    """AST class for break statement."""
+
+    def __init__(
+        self,
+        loc: SourceLocation = NO_SOURCE_LOCATION,
+        parent: Optional[ASTNodes] = None,
+    ) -> None:
+        """Initialize the BreakStmt instance."""
+        super().__init__(loc=loc, parent=parent)
+        self.kind = ASTKind.BreakStmtKind
+
+    def __str__(self) -> str:
+        """Return a string representation of the object."""
+        return "BreakStmt"
+
+    def get_struct(self, simplified: bool = False) -> ReprStruct:
+        """Return the AST structure of the object."""
+        key = "BREAK-STMT"
+        value: DictDataTypesStruct = {}
+        return self._prepare_struct(key, value, simplified)
+
+
+@public
+@typechecked
+class ContinueStmt(StatementType):
+    """AST class for continue statement."""
+
+    def __init__(
+        self,
+        loc: SourceLocation = NO_SOURCE_LOCATION,
+        parent: Optional[ASTNodes] = None,
+    ) -> None:
+        """Initialize the ContinueStmt instance."""
+        super().__init__(loc=loc, parent=parent)
+        self.kind = ASTKind.ContinueStmtKind
+
+    def __str__(self) -> str:
+        """Return a string representation of the object."""
+        return "ContinueStmt"
+
+    def get_struct(self, simplified: bool = False) -> ReprStruct:
+        """Return the AST structure of the object."""
+        key = "CONTINUE-STMT"
+        value: DictDataTypesStruct = {}
+        return self._prepare_struct(key, value, simplified)
+
+
 @public
 @typechecked
 class WhileStmt(StatementType):

libs/astx/tests/test_flows.py

@@ -1,14 +1,18 @@
 """Tests for control flow statements."""
 
+from typing import Any, cast
+
 import astx
 import pytest
 
-from astx.base import SourceLocation
+from astx.base import ReprStruct, SourceLocation
 from astx.blocks import Block
 from astx.flows import (
     AsyncForRangeLoopExpr,
     AsyncForRangeLoopStmt,
+    BreakStmt,
     CaseStmt,
+    ContinueStmt,
     DoWhileExpr,
     DoWhileStmt,
     ForCountLoopExpr,
@@ -159,6 +163,38 @@ def test_for_count_loop_expr() -> None:
     visualize(for_expr.get_struct())
 
 
+def test_break_stmt() -> None:
+    """Test BreakStmt class."""
+    break_stmt = BreakStmt()
+
+    assert str(break_stmt) == "BreakStmt"
+    assert break_stmt.kind == astx.ASTKind.BreakStmtKind
+
+    struct: ReprStruct = break_stmt.get_struct()
+    assert "BREAK-STMT" in cast(dict[str, Any], struct)
+
+    struct_simplified: ReprStruct = break_stmt.get_struct(simplified=True)
+    assert "BREAK-STMT" in cast(dict[str, Any], struct_simplified)
+
+    visualize(struct)
+
+
+def test_continue_stmt() -> None:
+    """Test ContinueStmt class."""
+    continue_stmt = ContinueStmt()
+
+    assert str(continue_stmt) == "ContinueStmt"
+    assert continue_stmt.kind == astx.ASTKind.ContinueStmtKind
+
+    struct: ReprStruct = continue_stmt.get_struct()
+    assert "CONTINUE-STMT" in cast(dict[str, Any], struct)
+
+    struct_simplified: ReprStruct = continue_stmt.get_struct(simplified=True)
+    assert "CONTINUE-STMT" in cast(dict[str, Any], struct_simplified)
+
+    visualize(struct)
+
+
 def test_async_for_range_loop_expr() -> None:
     """Test `Async For Range Loop` expression`."""
     decl_a = InlineVariableDeclaration(