feat: Add Dict Comprehension support (#198)

Author: Zehen-249

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

@@ -731,3 +731,15 @@ def visit(self, node: astx.GeneratorExpr) -> str:
             ret_str += f" if {self.visit(cond)}"
 
         return f"({ret_str})"
+
+    @dispatch  # type: ignore[no-redef]
+    def visit(self, node: astx.DictComprehension) -> str:
+        """Handle DictComprehension nodes."""
+        ret_str = (
+            f"{self.visit(node.key)}: {self.visit(node.value)}"
+            f" for {self.visit(node.target)} in {self.visit(node.iterable)}"
+        )
+        if hasattr(node, "conditions") and node.conditions:
+            for cond in node.conditions:
+                ret_str += f" if {self.visit(cond)}"
+        return f"{{{ret_str}}}"

libs/astx-transpilers/tests/test_python.py

@@ -1589,3 +1589,28 @@ def test_transpiler_generator_expr_no_conditions() -> None:
     assert generated_code == expected_code, (
         f"Expected '{expected_code}', but got '{generated_code}'"
     )
+
+
+def test_transpiler_dict_comprehension() -> None:
+    """Test astx.DictComprehension."""
+    dict_comp = astx.DictComprehension(
+        key=astx.Variable("x"),
+        value=astx.BinaryOp(
+            op_code="*", lhs=astx.Variable("x"), rhs=astx.Variable("x")
+        ),
+        target=astx.Variable("x"),
+        iterable=astx.Variable("range(10)"),
+        conditions=[
+            astx.BinaryOp(
+                op_code=">", lhs=astx.Variable("x"), rhs=astx.LiteralInt32(5)
+            ),
+            astx.BinaryOp(
+                op_code="<", lhs=astx.Variable("x"), rhs=astx.LiteralInt32(7)
+            ),
+        ],
+    )
+    generated_code = transpiler.visit(dict_comp)
+    expected_code = "{x: (x * x) for x in range(10) if (x > 5) if (x < 7)}"
+    assert generated_code == expected_code, (
+        f"Expected '{expected_code}', but got '{generated_code}'"
+    )

libs/astx/src/astx/__init__.py

@@ -64,6 +64,7 @@
     AsyncForRangeLoopExpr,
     AsyncForRangeLoopStmt,
     CaseStmt,
+    DictComprehension,
     DoWhileExpr,
     DoWhileStmt,
     ForCountLoopExpr,
@@ -234,6 +235,7 @@ def get_version() -> str:
     "Date",
     "DateTime",
     "DeleteStmt",
+    "DictComprehension",
     "DictType",
     "DoWhileExpr",
     "DoWhileStmt",

libs/astx/src/astx/base.py

@@ -138,6 +138,7 @@ class ASTKind(Enum):
     DoWhileStmtKind = -515
     DoWhileExprKind = -516
     GeneratorExprKind = -517
+    DictComprehensionKind = -518
 
     # data types
     NullDTKind = -600

libs/astx/src/astx/flows.py

@@ -18,6 +18,9 @@
     StatementType,
 )
 from astx.blocks import Block
+from astx.callables import (
+    Comprehension,
+)
 from astx.tools.typing import typechecked
 from astx.variables import InlineVariableDeclaration
 
@@ -792,3 +795,66 @@ def get_struct(self, simplified: bool = False) -> ReprStruct:
             else ASTNodes().get_struct(simplified),
         }
         return self._prepare_struct(key, value, simplified)
+
+
+@public
+@typechecked
+class DictComprehension(Comprehension):
+    """AST node for dictionary comprehensions."""
+
+    key: Expr
+    value: Expr
+
+    def __init__(
+        self,
+        key: Expr,
+        value: Expr,
+        target: Expr,
+        iterable: Expr,
+        conditions: Optional[Iterable[Expr] | ASTNodes[Expr]] = None,
+        is_async: bool = False,
+        loc: SourceLocation = NO_SOURCE_LOCATION,
+        parent: Optional[ASTNodes] = None,
+    ) -> None:
+        super().__init__(
+            target=target,
+            iterable=iterable,
+            conditions=conditions,
+            is_async=is_async,
+            loc=loc,
+            parent=parent,
+        )
+        self.key = key
+        self.value = value
+        self.kind = ASTKind.DictComprehensionKind
+
+    def __str__(self) -> str:
+        """Return a string representation of the object."""
+        conditions_str = []
+        if hasattr(self, "conditions") and self.conditions:
+            conditions_str = [str(cond) for cond in self.conditions]
+
+        ret_str = (
+            f"DictComprehension[key={self.key}, value={self.value},"
+            f" target={self.target}, iterable={self.iterable}, "
+            f" conditions={conditions_str},"
+            f" is_async={self.is_async}]"
+        )
+        return ret_str
+
+    def get_struct(self, simplified: bool = False) -> ReprStruct:
+        """Return the AST structure of the object."""
+        value: DictDataTypesStruct = {
+            "key": self.key.get_struct(simplified),
+            "value": self.value.get_struct(simplified),
+            "target": self.target.get_struct(simplified),
+            "iterable": self.iterable.get_struct(simplified),
+        }
+        if hasattr(self, "conditions") and self.conditions:
+            value["conditions"] = self.conditions.get_struct(simplified)
+        key = (
+            f"DICT-COMPREHENSION#{id(self)}"
+            if simplified
+            else "DICT-COMPREHENSION"
+        )
+        return self._prepare_struct(key, value, simplified)

libs/astx/tests/test_flows.py

@@ -9,6 +9,7 @@
     AsyncForRangeLoopExpr,
     AsyncForRangeLoopStmt,
     CaseStmt,
+    DictComprehension,
     DoWhileExpr,
     DoWhileStmt,
     ForCountLoopExpr,
@@ -437,3 +438,22 @@ def test_generator_expr_2() -> None:
     assert gen_expr.get_struct()
     assert gen_expr.get_struct(simplified=True)
     visualize(gen_expr.get_struct())
+
+
+def test_dict_comprehension() -> None:
+    """Test `DictComprehension` class."""
+    dict_comprehension = DictComprehension(
+        key=Variable("x"),
+        value=BinaryOp(op_code="*", lhs=Variable("x"), rhs=Variable("x")),
+        target=Variable("x"),
+        iterable=Variable("range(10)"),
+        conditions=[
+            BinaryOp(op_code=">", lhs=Variable("x"), rhs=LiteralInt32(5)),
+            BinaryOp(op_code="<", lhs=Variable("x"), rhs=LiteralInt32(7)),
+        ],
+    )
+
+    assert str(dict_comprehension)
+    assert dict_comprehension.get_struct()
+    assert dict_comprehension.get_struct(simplified=True)
+    visualize(dict_comprehension.get_struct())