Add dict comprehension Support

docs/tutorials/for-loop.ipynb

@@ -28,7 +28,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "id": "45319c53",
    "metadata": {},
    "outputs": [],

src/astx/__init__.py

@@ -62,6 +62,7 @@
     AsyncForRangeLoopExpr,
     AsyncForRangeLoopStmt,
     CaseStmt,
+    DictComprehension,
     ForCountLoopExpr,
     ForCountLoopStmt,
     ForRangeLoopExpr,
@@ -224,6 +225,7 @@ def get_version() -> str:
     "DataTypeOps",
     "Date",
     "DateTime",
+    "DictComprehension",
     "DictType",
     "EnumDeclStmt",
     "ExceptionHandlerStmt",

src/astx/base.py

@@ -131,6 +131,7 @@ class ASTKind(Enum):
     WithStmtKind = -512
     AsyncRangeLoopStmtKind = -513
     AsyncRangeLoopExprKind = -514
+    DictComprehensionKind = -515
 
     # data types
     NullDTKind = -600

src/astx/flows.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Optional, cast
+from typing import Optional, Union, cast
 
 from public import public
 
@@ -18,6 +18,7 @@
     StatementType,
 )
 from astx.blocks import Block
+from astx.literals import LiteralList, LiteralSet, LiteralString, LiteralTuple
 from astx.tools.typing import typechecked
 from astx.variables import InlineVariableDeclaration
 
@@ -687,3 +688,85 @@ def get_struct(self, simplified: bool = False) -> ReprStruct:
         key = f"GOTO-STMT[{self.label.value}]"
         value: DictDataTypesStruct = {}
         return self._prepare_struct(key, value, simplified)
+
+
+@public
+@typechecked
+class DictComprehension(Expr):
+    """AST class for function `DictComprehension` statement."""
+
+    def __init__(
+        self,
+        key: Identifier,
+        value: Identifier,
+        iterable: Identifier,
+        iterator: Union[
+            LiteralList, LiteralTuple, LiteralSet, LiteralString, Identifier
+        ],
+        loc: SourceLocation = NO_SOURCE_LOCATION,
+        parent: Optional[ASTNodes] = None,
+    ) -> None:
+        """Initialize the Return instance."""
+        super().__init__(loc, parent)
+        self.key = key
+        self.value = value
+        self.iterable = iterable
+        self.iterator = iterator
+        self.kind = ASTKind.DictComprehensionKind
+
+    def __str__(self) -> str:
+        """Return a string representation of the object."""
+        if isinstance(self.iterator, LiteralList):
+            elements = ", ".join(str(e.value) for e in self.iterator.elements)
+            return (
+                f"{{{self.key.value}: {self.value.value} for "
+                f"{self.iterable.value} in [{elements}]}}"
+            )
+        elif isinstance(self.iterator, LiteralTuple):
+            elements = ", ".join(str(e.value) for e in self.iterator.elements)
+            return (
+                f"{{{self.key.value}: {self.value.value} for "
+                f"{self.iterable.value} in ({elements})}}"
+            )
+        elif isinstance(self.iterator, LiteralSet):
+            elements = ", ".join(str(e.value) for e in self.iterator.elements)
+            return (
+                f"{{{self.key.value}: {self.value.value} for "
+                f"{self.iterable.value} in {{{elements}}}}}"
+            )
+        elif isinstance(self.iterator, Identifier):
+            return (
+                f"{{{self.key.value}: {self.value.value} for "
+                f"{self.iterable.value} in {self.iterator.value}}}"
+            )
+        else:
+            return (
+                f"{{{self.key.value}: {self.value.value} for "
+                f"{self.iterable.value} in {self.iterator}}}"
+            )
+
+    def get_struct(self, simplified: bool = False) -> ReprStruct:
+        """Return the AST structure of the object."""
+        key = "DICT-COMPREHENSION"
+        value: DictDataTypesStruct = {
+            "key": self.key.get_struct(simplified),
+            "value": self.value.get_struct(simplified),
+            "iterable": self.iterable.get_struct(simplified),
+        }
+
+        if isinstance(self.iterator, (LiteralList, LiteralSet, LiteralTuple)):
+            elements_struct = [
+                element.get_struct(simplified)
+                for element in self.iterator.elements
+            ]
+            iterator_struct = {
+                f"{self.iterator.__class__.__name__}": {
+                    f"ELEMENT[{i}]": element
+                    for i, element in enumerate(elements_struct)
+                }
+            }
+            value["iterator"] = cast(DictDataTypesStruct, iterator_struct)
+        else:
+            value["iterator"] = self.iterator.get_struct(simplified)
+
+        return self._prepare_struct(key, value, simplified)

src/astx/tools/transpilers/python.py

@@ -682,3 +682,13 @@ def visit(self, node: astx.LiteralDict) -> str:
             for key, value in node.elements.items()
         )
         return f"{{{items_code}}}"
+
+    @dispatch  # type: ignore[no-redef]
+    def visit(self, node: astx.DictComprehension) -> str:
+        """Handle DictComprehension nodes."""
+        key = self.visit(node.key)
+        value = self.visit(node.value)
+        iterable = self.visit(node.iterable)
+        iterator = self.visit(node.iterator)
+
+        return f"{{{key}: {value} for {iterable} in {iterator}}}"

tests/test_flows.py

@@ -1,5 +1,10 @@
 """Tests for control flow statements."""
 
+# for running tests from local astx module
+# import os,sys
+# # print(os.path.abspath(os.path.join(os.getcwd(), ".", "src")))
+# sys.path.insert(0, os.path.abspath(os.path.join(os.getcwd(), ".", "src")))
+
 import astx
 import pytest
 
@@ -323,3 +328,17 @@ def test_goto_stmt() -> None:
     assert goto_stmt.get_struct()
     assert goto_stmt.get_struct(simplified=True)
     visualize(goto_stmt.get_struct())
+
+
+def test_dict_comprehension() -> None:
+    """Test `DictComprehension` class."""
+    dict_comp = astx.DictComprehension(
+        key=astx.Identifier("x"),
+        value=astx.Identifier("x*x"),
+        iterable=astx.Identifier("x"),
+        iterator=astx.Identifier("my_list"),
+    )
+    assert str(dict_comp) == "{x: x*x for x in my_list}"
+    assert dict_comp.get_struct()
+    assert dict_comp.get_struct(simplified=True)
+    visualize(dict_comp.get_struct())

tests/tools/transpilers/test_python.py

@@ -1,5 +1,10 @@
 """Test Python Transpiler."""
 
+# for running tests from local astx module
+# import os,sys
+# # print(os.path.abspath(os.path.join(os.getcwd(), ".", "src")))
+# sys.path.insert(0, os.path.abspath(os.path.join(os.getcwd(), ".", "src")))
+
 import ast
 import sys
 
@@ -1397,3 +1402,40 @@ def test_transpiler_literal_dict() -> None:
         1: 10,
         2: 20,
     }, f"Expected '{expected_code}', but got '{generated_code}'"
+
+
+def test_dict_comprehension() -> None:
+    """Test astx.DictComprehension."""
+    dict_comp_1 = astx.DictComprehension(
+        key=astx.Identifier("x"),
+        value=astx.Identifier("x*x"),
+        iterable=astx.Identifier("x"),
+        iterator=astx.Identifier("my_list"),
+    )
+
+    generated_code_1 = transpiler.visit(dict_comp_1)
+    expected_code_1 = "{x: x*x for x in my_list}"
+
+    dict_comp_2 = astx.DictComprehension(
+        key=astx.Identifier("x"),
+        value=astx.Identifier("x*x"),
+        iterable=astx.Identifier("x"),
+        iterator=astx.LiteralList(
+            elements=[
+                astx.LiteralInt32(10),
+                astx.LiteralInt32(20),
+                astx.LiteralInt32(30),
+            ]
+        ),
+    )
+
+    generated_code_2 = transpiler.visit(dict_comp_2)
+    expected_code_2 = "{x: x*x for x in [10, 20, 30]}"
+
+    assert generated_code_1 == expected_code_1, (
+        f"Expected code: {expected_code_1} ;"
+        f" Generated code: {generated_code_1}"
+    )
+    assert generated_code_2 == expected_code_2, (
+        f"Expected code: {expected_code_2} ;Generated code: {generated_code_2}"
+    )