diff --git a/docs/tutorials/for-loop.ipynb b/docs/tutorials/for-loop.ipynb index 8859621..824bd52 100644 --- a/docs/tutorials/for-loop.ipynb +++ b/docs/tutorials/for-loop.ipynb @@ -148,7 +148,7 @@ "body_for = astx.Block()\n", "\n", "# Define a For Count Loop\n", - "for_counter = astx.ForCountLoop(\n", + "for_counter = astx.ForCountLoopStmt(\n", " initializer=decl_a,\n", " condition=var_a < astx.LiteralInt32(10),\n", " update=astx.UnaryOp(\"++\", var_a),\n", @@ -161,7 +161,7 @@ ], "metadata": { "kernelspec": { - "display_name": "astx", + "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, @@ -175,7 +175,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.12.3" + "version": "3.12.7" } }, "nbformat": 4, diff --git a/src/astx/__init__.py b/src/astx/__init__.py index 8fd47ff..cf7f7c0 100644 --- a/src/astx/__init__.py +++ b/src/astx/__init__.py @@ -85,7 +85,8 @@ UTF8String, ) from astx.flows import ( - ForCountLoop, + ForCountLoopExpr, + ForCountLoopStmt, ForRangeLoopExpr, ForRangeLoopStmt, If, @@ -158,7 +159,8 @@ def get_version() -> str: "LiteralFloat64", "Floating", "flows", - "ForCountLoop", + "ForCountLoopStmt", + "ForCountLoopExpr", "ForRangeLoopStmt", "ForRangeLoopExpr", "Function", diff --git a/src/astx/base.py b/src/astx/base.py index 1c383e6..733befe 100644 --- a/src/astx/base.py +++ b/src/astx/base.py @@ -97,10 +97,11 @@ class ASTKind(Enum): # control flow IfKind = -500 - ForCountKind = -501 + ForCountLoopStmtKind = -501 ForRangeLoopStmtKind = -502 WhileKind = -503 ForRangeLoopExprKind = -504 + ForCountLoopExprKind = -505 # data types NullDTKind = -600 diff --git a/src/astx/flows.py b/src/astx/flows.py index 11639e4..823cf6d 100644 --- a/src/astx/flows.py +++ b/src/astx/flows.py @@ -185,7 +185,7 @@ def get_struct(self, simplified: bool = False) -> ReprStruct: @public @typechecked -class ForCountLoop(StatementType): +class ForCountLoopStmt(StatementType): """ AST class for a simple Count-Controlled `For` Loop statement. @@ -206,20 +206,20 @@ def __init__( loc: SourceLocation = NO_SOURCE_LOCATION, parent: Optional[ASTNodes] = None, ) -> None: - """Initialize the ForStmt instance.""" + """Initialize the ForCountLoopStmt instance.""" super().__init__(loc=loc, parent=parent) self.initializer = initializer self.condition = condition self.update = update self.body = body - self.kind = ASTKind.ForCountKind + self.kind = ASTKind.ForCountLoopStmtKind def __str__(self) -> str: """Return a string that represents the object.""" init = self.initializer cond = self.condition update = self.update - return f"ForCountLoop({init};{cond};{update})" + return f"ForCountLoopStmt({init};{cond};{update})" def get_struct(self, simplified: bool = False) -> ReprStruct: """Return the AST structure of the object.""" @@ -239,6 +239,61 @@ def get_struct(self, simplified: bool = False) -> ReprStruct: return self._prepare_struct(key, value, simplified) +@public +@typechecked +class ForCountLoopExpr(Expr): + """ + AST class for a simple Count-Controlled `For` Loop expression. + + This is a very basic `for` loop, used by languages like C or C++. + """ + + initializer: InlineVariableDeclaration + condition: Expr + update: Expr + body: Block + + def __init__( + self, + initializer: InlineVariableDeclaration, + condition: Expr, + update: Expr, + body: Block, + loc: SourceLocation = NO_SOURCE_LOCATION, + parent: Optional[ASTNodes] = None, + ) -> None: + """Initialize the ForLoopCountExpr instance.""" + super().__init__(loc=loc, parent=parent) + self.initializer = initializer + self.condition = condition + self.update = update + self.body = body + self.kind = ASTKind.ForCountLoopExprKind + + def __str__(self) -> str: + """Return a string that represents the object.""" + init = self.initializer + cond = self.condition + update = self.update + return f"ForCountLoopExpr({init};{cond};{update})" + + def get_struct(self, simplified: bool = False) -> ReprStruct: + """Return the AST structure of the object.""" + for_init = {"initialization": self.initializer.get_struct(simplified)} + for_cond = {"condition": self.condition.get_struct(simplified)} + for_update = {"update": self.update.get_struct(simplified)} + for_body = self.body.get_struct(simplified) + + key = "FOR-COUNT-EXPR" + value: ReprStruct = { + **cast(DictDataTypesStruct, for_init), + **cast(DictDataTypesStruct, for_cond), + **cast(DictDataTypesStruct, for_update), + **cast(DictDataTypesStruct, for_body), + } + return self._prepare_struct(key, value, simplified) + + @public @typechecked class While(StatementType): diff --git a/tests/test_flows.py b/tests/test_flows.py index 0129e98..8e555f4 100644 --- a/tests/test_flows.py +++ b/tests/test_flows.py @@ -2,7 +2,13 @@ from astx.blocks import Block from astx.datatypes import Int32, LiteralInt32 -from astx.flows import ForCountLoop, ForRangeLoopExpr, ForRangeLoopStmt, If +from astx.flows import ( + ForCountLoopExpr, + ForCountLoopStmt, + ForRangeLoopExpr, + ForRangeLoopStmt, + If, +) from astx.operators import BinaryOp, UnaryOp from astx.variables import InlineVariableDeclaration, Variable from astx.viz import visualize @@ -73,7 +79,7 @@ def test_for_range_loop_stmt() -> None: visualize(for_stmt.get_struct()) -def test_for_count() -> None: +def test_for_count_loop_stmt() -> None: """Test `For Count Loop` statement.""" decl_a = InlineVariableDeclaration("a", type_=Int32, value=LiteralInt32(0)) var_a = Variable("a") @@ -81,7 +87,7 @@ def test_for_count() -> None: update = UnaryOp(op_code="++", operand=var_a) body = Block() body.append(LiteralInt32(2)) - for_stmt = ForCountLoop( + for_stmt = ForCountLoopStmt( initializer=decl_a, condition=cond, update=update, body=body ) @@ -89,3 +95,21 @@ def test_for_count() -> None: assert for_stmt.get_struct() assert for_stmt.get_struct(simplified=True) visualize(for_stmt.get_struct()) + + +def test_for_count_loop_expr() -> None: + """Test `For Count Loop` expression.""" + decl_a = InlineVariableDeclaration("a", type_=Int32, value=LiteralInt32(0)) + var_a = Variable("a") + cond = BinaryOp(op_code="<", lhs=var_a, rhs=LiteralInt32(10)) + update = UnaryOp(op_code="++", operand=var_a) + body = Block() + body.append(LiteralInt32(2)) + for_expr = ForCountLoopExpr( + initializer=decl_a, condition=cond, update=update, body=body + ) + + assert str(for_expr) + assert for_expr.get_struct() + assert for_expr.get_struct(simplified=True) + visualize(for_expr.get_struct())