diff --git a/libs/astx-transpilers/src/astx_transpilers/python_string.py b/libs/astx-transpilers/src/astx_transpilers/python_string.py index 51906894..0c9346cc 100644 --- a/libs/astx-transpilers/src/astx_transpilers/python_string.py +++ b/libs/astx-transpilers/src/astx_transpilers/python_string.py @@ -713,3 +713,15 @@ def visit(self, node: astx.DoWhileStmt) -> str: body = self._generate_block(node.body) condition = self.visit(node.condition) return f"while True:\n{body}\n if not {condition}:\n break" + + @dispatch # type: ignore[no-redef] + def visit(self, node: astx.GeneratorExpr) -> str: + """Handle GeneratorExpr nodes.""" + ret_str = ( + f"{self.visit(node.element)} for {self.visit(node.target)}" + f" in {self.visit(node.iterable)}" + ) + for cond in node.conditions: + ret_str += f" if {self.visit(cond)}" + + return f"({ret_str})" diff --git a/libs/astx-transpilers/tests/test_python.py b/libs/astx-transpilers/tests/test_python.py index 35bce373..491ba263 100644 --- a/libs/astx-transpilers/tests/test_python.py +++ b/libs/astx-transpilers/tests/test_python.py @@ -1525,3 +1525,44 @@ def test_transpiler_do_while_expr() -> None: assert generated_code == expected_code, ( f"Expected '{expected_code}', but got '{generated_code}'" ) + + +def test_transpiler_generator_expr() -> None: + """Test astx.GeneratorExpr.""" + 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=[ + astx.BinaryOp( + op_code=">", lhs=astx.Variable("x"), rhs=astx.LiteralInt32(3) + ), + astx.BinaryOp( + op_code="<", lhs=astx.Variable("x"), rhs=astx.LiteralInt32(7) + ), + ], + ) + 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, ( + f"Expected '{expected_code}', but got '{generated_code}'" + ) + + +def test_transpiler_generator_expr_no_conditions() -> None: + """Test astx.GeneratorExpr with no conditions.""" + gen_expr = astx.GeneratorExpr( + target=astx.Variable("x"), + element=astx.BinaryOp( + op_code="+", lhs=astx.Variable("x"), rhs=astx.Variable("x") + ), + iterable=astx.Identifier("range(10)"), + ) + + generated_code = translate(gen_expr) + expected_code = "((x + x) for x in range(10))" + assert generated_code == expected_code, ( + f"Expected '{expected_code}', but got '{generated_code}'" + ) diff --git a/libs/astx/src/astx/__init__.py b/libs/astx/src/astx/__init__.py index 4cd7709a..a53d4130 100644 --- a/libs/astx/src/astx/__init__.py +++ b/libs/astx/src/astx/__init__.py @@ -69,6 +69,7 @@ ForCountLoopStmt, ForRangeLoopExpr, ForRangeLoopStmt, + GeneratorExpr, GotoStmt, IfExpr, IfStmt, @@ -253,6 +254,7 @@ def get_version() -> str: "FunctionDef", "FunctionPrototype", "FunctionReturn", + "GeneratorExpr", "GotoStmt", "Identifier", "IfExpr", diff --git a/libs/astx/src/astx/base.py b/libs/astx/src/astx/base.py index 83ba78ad..228c070d 100644 --- a/libs/astx/src/astx/base.py +++ b/libs/astx/src/astx/base.py @@ -136,6 +136,7 @@ class ASTKind(Enum): AsyncRangeLoopExprKind = -514 DoWhileStmtKind = -515 DoWhileExprKind = -516 + GeneratorExprKind = -517 # data types NullDTKind = -600 diff --git a/libs/astx/src/astx/flows.py b/libs/astx/src/astx/flows.py index c95267a4..a87bce03 100644 --- a/libs/astx/src/astx/flows.py +++ b/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 @@ -737,3 +737,58 @@ def __init__( def __str__(self) -> str: """Return a string representation of the object.""" return f"DoWhileExpr[{self.condition}]" + + +@public +@typechecked +class GeneratorExpr(Expr): + """AST class for generator expressions.""" + + element: Expr + target: Expr + iterable: Expr + conditions: ASTNodes[Expr] | Iterable[Expr] + + def __init__( + self, + element: Expr, + target: Expr, + iterable: Expr, + conditions: Optional[ASTNodes[Expr]] | Optional[Iterable[Expr]] = [], + loc: SourceLocation = NO_SOURCE_LOCATION, + parent: Optional[ASTNodes] = None, + ) -> None: + """Initialize the GeneratorExpr instance.""" + super().__init__(loc=loc, parent=parent) + self.element = element + self.target = target + self.iterable = iterable + 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}, conditions=" + f"{[str(cond) for cond in self.conditions]}" + ) + return ret_str + + def get_struct(self, simplified: bool = False) -> ReprStruct: + """Return the AST structure of the object.""" + key = f"GENERATOR-EXPR#{id(self)}" if simplified else "GENERATOR-EXPR" + value: ReprStruct = { + "element": self.element.get_struct(simplified), + "target": self.target.get_struct(simplified), + "iterable": self.iterable.get_struct(simplified), + "conditions": self.conditions.get_struct(simplified) + if isinstance(self.conditions, ASTNodes) + else ASTNodes().get_struct(simplified), + } + return self._prepare_struct(key, value, simplified) diff --git a/libs/astx/tests/test_flows.py b/libs/astx/tests/test_flows.py index 45f97765..e31dd26a 100644 --- a/libs/astx/tests/test_flows.py +++ b/libs/astx/tests/test_flows.py @@ -386,3 +386,54 @@ def test_do_while_stmt() -> None: assert do_while_stmt.get_struct() assert do_while_stmt.get_struct(simplified=True) visualize(do_while_stmt.get_struct()) + + +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") + ), + target=astx.Variable("x"), + iterable=astx.Identifier("range(10)"), + conditions=[ + astx.BinaryOp( + op_code="<", lhs=astx.Variable("x"), rhs=astx.LiteralInt32(8) + ), + astx.BinaryOp( + 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())