fix: update conditions paramenter type

Author: Zehen-249

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

@@ -721,7 +721,7 @@ def visit(self, node: astx.GeneratorExpr) -> str:
             f"{self.visit(node.element)} for {self.visit(node.target)}"
             f" in {self.visit(node.iterable)}"
         )
-        if node.conditions:
-            for condition in node.conditions:
-                ret_str += f" if {self.visit(condition)}"
+        for cond in node.conditions:
+            ret_str += f" if {self.visit(cond)}"
+
         return f"({ret_str})"

libs/astx-transpilers/tests/test_python.py

@@ -1544,7 +1544,6 @@ def test_transpiler_generator_expr() -> None:
             ),
         ],
     )
-    print
     generated_code = translate(gen_expr)
     expected_code = "((x + x) for x in range(10) if (x > 3) if (x < 7))"
     assert generated_code == expected_code, (

libs/astx/src/astx/flows.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Optional, cast
+from typing import Iterable, Optional, cast
 
 from public import public
 
@@ -747,14 +747,14 @@ class GeneratorExpr(Expr):
     element: Expr
     target: Expr
     iterable: Expr
-    conditions: list[Expr]
+    conditions: ASTNodes[Expr] | Iterable[Expr]
 
     def __init__(
         self,
         element: Expr,
         target: Expr,
         iterable: Expr,
-        conditions: Optional[list[Expr]] = None,
+        conditions: Optional[ASTNodes[Expr]] | Optional[Iterable[Expr]] = [],
         loc: SourceLocation = NO_SOURCE_LOCATION,
         parent: Optional[ASTNodes] = None,
     ) -> None:
@@ -763,20 +763,21 @@ def __init__(
         self.element = element
         self.target = target
         self.iterable = iterable
-        self.conditions = conditions if conditions is not None else []
+        if isinstance(conditions, ASTNodes):
+            self.conditions = conditions
+        elif isinstance(conditions, Iterable):
+            self.conditions = ASTNodes()
+            for condition in conditions:
+                self.conditions.append(condition)
         self.kind = ASTKind.GeneratorExprKind
 
     def __str__(self) -> str:
         """Return a string representation of the object."""
         ret_str = (
             f"GeneratorExpr[element={self.element}, target={self.target},"
-            f" iterable={self.iterable},"
+            f" iterable={self.iterable}, conditions="
+            f"{[str(cond) for cond in self.conditions]}"
         )
-        if self.conditions is not None:
-            cons_list = []
-            for cond in self.conditions:
-                cons_list.append(str(cond))
-            ret_str += f" conditions={(cons_list)}]"
         return ret_str
 
     def get_struct(self, simplified: bool = False) -> ReprStruct:
@@ -786,12 +787,8 @@ def get_struct(self, simplified: bool = False) -> ReprStruct:
             "element": self.element.get_struct(simplified),
             "target": self.target.get_struct(simplified),
             "iterable": self.iterable.get_struct(simplified),
-            "conditions": cast(
-                ReprStruct,
-                {
-                    str(cond): cond.get_struct(simplified)
-                    for cond in self.conditions
-                },
-            ),
+            "conditions": self.conditions.get_struct(simplified)
+            if isinstance(self.conditions, ASTNodes)
+            else ASTNodes().get_struct(simplified),
         }
         return self._prepare_struct(key, value, simplified)

libs/astx/tests/test_flows.py

@@ -388,8 +388,8 @@ def test_do_while_stmt() -> None:
     visualize(do_while_stmt.get_struct())
 
 
-def test_generator_expr() -> None:
-    """Test `GeneratorExpr` class."""
+def test_generator_expr_1() -> None:
+    """Test `GeneratorExpr` class with conditions of Iterable type."""
     gen_expr = astx.GeneratorExpr(
         element=astx.BinaryOp(
             op_code="+", lhs=astx.Variable("x"), rhs=astx.Variable("x")
@@ -398,14 +398,42 @@ def test_generator_expr() -> None:
         iterable=astx.Identifier("range(10)"),
         conditions=[
             astx.BinaryOp(
-                op_code=">", lhs=astx.Variable("x"), rhs=astx.LiteralInt32(3)
+                op_code="<", lhs=astx.Variable("x"), rhs=astx.LiteralInt32(8)
             ),
             astx.BinaryOp(
-                op_code="<", lhs=astx.Variable("x"), rhs=astx.LiteralInt32(7)
+                op_code="%", lhs=astx.Variable("x"), rhs=astx.LiteralInt32(2)
             ),
         ],
     )
     assert str(gen_expr)
     assert gen_expr.get_struct()
     assert gen_expr.get_struct(simplified=True)
     visualize(gen_expr.get_struct())
+
+
+def test_generator_expr_2() -> None:
+    """Test `GeneratorExpr` class with conditions of ASTNodes type."""
+    conditions = astx.ASTNodes[astx.Expr]()
+    conditions.append(
+        astx.BinaryOp(
+            op_code="<", lhs=astx.Variable("x"), rhs=astx.LiteralInt32(8)
+        )
+    )
+    conditions.append(
+        astx.BinaryOp(
+            op_code="%", lhs=astx.Variable("x"), rhs=astx.LiteralInt32(2)
+        )
+    )
+    gen_expr = astx.GeneratorExpr(
+        element=astx.BinaryOp(
+            op_code="+", lhs=astx.Variable("x"), rhs=astx.Variable("x")
+        ),
+        target=astx.Variable("x"),
+        iterable=astx.Identifier("range(10)"),
+        conditions=conditions,
+    )
+    print(gen_expr)
+    assert str(gen_expr)
+    assert gen_expr.get_struct()
+    assert gen_expr.get_struct(simplified=True)
+    visualize(gen_expr.get_struct())