feat: add ForCountLoopExpr class (#141)

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,

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",

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

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):

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,19 +79,37 @@ 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")
     cond = BinaryOp(op_code="<", lhs=var_a, rhs=LiteralInt32(10))
     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
     )
 
     assert str(for_stmt)
     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())