Generatorexp support

docs/tutorials/context.ipynb

@@ -229,7 +229,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "ast",
    "language": "python",
    "name": "python3"
   },
@@ -243,7 +243,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.7"
+   "version": "3.13.2"
   }
  },
  "nbformat": 4,

src/astx/__init__.py

@@ -43,6 +43,7 @@
     FunctionReturn,
     LambdaExpr,
     YieldExpr,
+    YieldFromExpr,
 )
 from astx.classes import (
     ClassDeclStmt,
@@ -58,12 +59,15 @@
     ThrowStmt,
 )
 from astx.flows import (
+    AsyncForRangeLoopExpr,
+    AsyncForRangeLoopStmt,
     CaseStmt,
     ForCountLoopExpr,
     ForCountLoopStmt,
     ForRangeLoopExpr,
     ForRangeLoopStmt,
     GotoStmt,
+    GeneratorExpr,
     IfExpr,
     IfStmt,
     SwitchStmt,
@@ -199,6 +203,8 @@ def get_version() -> str:
     "Argument",
     "Arguments",
     "AssignmentExpr",
+    "AsyncForRangeLoopExpr",
+    "AsyncForRangeLoopStmt",
     "AwaitExpr",
     "BinaryOp",
     "Block",
@@ -236,6 +242,7 @@ def get_version() -> str:
     "FunctionDef",
     "FunctionPrototype",
     "FunctionReturn",
+    "GeneratorExpr",
     "GotoStmt",
     "Identifier",
     "IfExpr",
@@ -330,6 +337,7 @@ def get_version() -> str:
     "XnorOp",
     "XorOp",
     "YieldExpr",
+    "YieldFromExpr",
     "base",
     "blocks",
     "callables",

src/astx/base.py

@@ -113,6 +113,7 @@ class ASTKind(Enum):
     FunctionAsyncDefKind = -405
     AwaitExprKind = -406
     YieldExprKind = -510
+    YieldFromExprKind = -407
 
     # control flow
     IfStmtKind = -500
@@ -127,6 +128,9 @@ class ASTKind(Enum):
     SwitchStmtKind = -509
     GotoStmtKind = -511
     WithStmtKind = -512
+    GeneratorExprKind = -516
+    AsyncRangeLoopStmtKind = -513
+    AsyncRangeLoopExprKind = -514
 
     # data types
     NullDTKind = -600

src/astx/callables.py

@@ -386,3 +386,32 @@ def get_struct(self, simplified: bool = False) -> ReprStruct:
         key = "YIELD-EXPR"
         value = {} if self.value is None else self.value.get_struct(simplified)
         return self._prepare_struct(key, value, simplified)
+
+
+@public
+@typechecked
+class YieldFromExpr(Expr):
+    """AST class for YieldFromExpr."""
+
+    value: Expr
+
+    def __init__(
+        self,
+        value: Expr,
+        loc: SourceLocation = NO_SOURCE_LOCATION,
+        parent: Optional[ASTNodes] = None,
+    ) -> None:
+        """Initialize the YieldFromExpr instance."""
+        super().__init__(loc=loc, parent=parent)
+        self.value = value
+        self.kind = ASTKind.YieldFromExprKind
+
+    def __str__(self) -> str:
+        """Return a string representation of the object."""
+        return f"YieldFromExpr[{self.value}]"
+
+    def get_struct(self, simplified: bool = False) -> ReprStruct:
+        """Return the AST structure of the object."""
+        key = "YIELDFROM-EXPR"
+        value = self.value.get_struct(simplified)
+        return self._prepare_struct(key, value, simplified)

src/astx/flows.py

@@ -2,7 +2,8 @@
 
 from __future__ import annotations
 
-from typing import Optional, cast
+from typing import Optional, cast, Union
+
 
 from public import public
 
@@ -20,7 +21,7 @@
 from astx.blocks import Block
 from astx.tools.typing import typechecked
 from astx.variables import InlineVariableDeclaration
-
+from astx.literals import LiteralString, LiteralList, LiteralTuple, LiteralSet
 
 @public
 @typechecked
@@ -343,6 +344,131 @@ def get_struct(self, simplified: bool = False) -> ReprStruct:
         return self._prepare_struct(key, value, simplified)
 
 
+@public
+@typechecked
+class AsyncForRangeLoopStmt(StatementType):
+    """AST class for asynchronous `For` Range Statement."""
+
+    variable: InlineVariableDeclaration
+    start: Optional[Expr]
+    end: Expr
+    step: Optional[Expr]
+    body: Block
+
+    def __init__(
+        self,
+        variable: InlineVariableDeclaration,
+        start: Optional[Expr],
+        end: Expr,
+        step: Optional[Expr],
+        body: Block,
+        loc: SourceLocation = NO_SOURCE_LOCATION,
+        parent: Optional[ASTNodes] = None,
+    ) -> None:
+        """Initialize the AsyncForRangeLoopStmt instance."""
+        super().__init__(loc=loc, parent=parent)
+        self.variable = variable
+        self.start = start
+        self.end = end
+        self.step = step
+        self.body = body
+        self.kind = ASTKind.AsyncRangeLoopStmtKind
+
+    def __str__(self) -> str:
+        """Return a string that represents the object."""
+        start = self.start
+        end = self.end
+        step = self.step
+        var_name = self.variable.name
+        return f"AsyncForRangeLoopStmt({var_name}=[{start}:{end}:{step}])"
+
+    def get_struct(self, simplified: bool = False) -> ReprStruct:
+        """Return the AST structure of the object."""
+        for_start = {
+            "start": {}
+            if self.start is None
+            else self.start.get_struct(simplified)
+        }
+        for_end = {"end": self.end.get_struct(simplified)}
+        for_step = {
+            "step": {}
+            if self.step is None
+            else self.step.get_struct(simplified)
+        }
+        for_body = self.body.get_struct(simplified)
+
+        key = "ASYNC-FOR-RANGE-LOOP-STMT"
+        value: ReprStruct = {
+            **cast(DictDataTypesStruct, for_start),
+            **cast(DictDataTypesStruct, for_end),
+            **cast(DictDataTypesStruct, for_step),
+            **cast(DictDataTypesStruct, for_body),
+        }
+        return self._prepare_struct(key, value, simplified)
+
+
+@public
+@typechecked
+class AsyncForRangeLoopExpr(Expr):
+    """AST class for asynchronous `For` Range Expression."""
+
+    variable: InlineVariableDeclaration
+    start: Optional[Expr]
+    end: Expr
+    step: Optional[Expr]
+    body: Block
+
+    def __init__(
+        self,
+        variable: InlineVariableDeclaration,
+        start: Optional[Expr],
+        end: Expr,
+        step: Optional[Expr],
+        body: Block,
+        loc: SourceLocation = NO_SOURCE_LOCATION,
+        parent: Optional[ASTNodes] = None,
+    ) -> None:
+        """Initialize the AsyncForRangeLoopExpr instance."""
+        super().__init__(loc=loc, parent=parent)
+        self.variable = variable
+        self.start = start
+        self.end = end
+        self.step = step
+        self.body = body
+        self.kind = ASTKind.AsyncRangeLoopExprKind
+
+    def __str__(self) -> str:
+        """Return a string that represents the object."""
+        var_name = self.variable.name
+        return f"AsyncForRangeLoopExpr[{var_name}]"
+
+    def get_struct(self, simplified: bool = False) -> ReprStruct:
+        """Return the AST structure of the object."""
+        for_var = {"var": self.variable.get_struct(simplified)}
+        for_start = {
+            "start": {}
+            if self.start is None
+            else self.start.get_struct(simplified)
+        }
+        for_end = {"end": self.end.get_struct(simplified)}
+        for_step = {
+            "step": {}
+            if self.step is None
+            else self.step.get_struct(simplified)
+        }
+        for_body = self.body.get_struct(simplified)
+
+        key = "ASYNC-FOR-RANGE-LOOP-EXPR"
+        value: ReprStruct = {
+            **cast(DictDataTypesStruct, for_var),
+            **cast(DictDataTypesStruct, for_start),
+            **cast(DictDataTypesStruct, for_end),
+            **cast(DictDataTypesStruct, for_step),
+            **cast(DictDataTypesStruct, for_body),
+        }
+        return self._prepare_struct(key, value, simplified)
+
+
 @public
 @typechecked
 class WhileStmt(StatementType):
@@ -562,3 +688,55 @@ 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 GeneratorExpr(Expr):
+    """AST class for generator expressions."""
+
+    element: Identifier
+    target: Identifier
+    iterator:  Union[LiteralList, LiteralTuple, LiteralSet, LiteralString, Identifier]
+    def __init__(
+        self,
+        element: Identifier,
+        target: Identifier,
+        iterator:  Union[LiteralList, LiteralTuple, LiteralSet, LiteralString, Identifier],
+        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.iterator = iterator
+        self.kind = ASTKind.GeneratorExprKind
+
+    def __str__(self) -> str:
+        """Return a string representation of the object."""
+        if isinstance(self.iterator, LiteralList):
+            return f"({self.element.value} for {self.target.value} in [{", ".join(str(e.value) for e in self.iterator.elements)}])"
+        elif isinstance(self.iterator, LiteralTuple):
+            return f"({self.element.value} for {self.target.value} in ({", ".join(str(e.value) for e in self.iterator.elements)}))"
+        elif isinstance(self.iterator, LiteralSet):
+            return f"({self.element.value} for {self.target.value} in {{{", ".join(str(e.value) for e in self.iterator.elements)}}})"
+        elif isinstance(self.iterator, Identifier):
+            return f"({self.element.value} for {self.target.value} in {self.iterator.value})"
+        else:
+            return f"({self.element.value} for {self.target.value} in {self.iterator})"
+
+    def get_struct(self, simplified: bool = False) -> ReprStruct:
+        """Return the AST structure of the object."""
+        key = "GENERATOR-EXPR"
+        value: ReprStruct = {
+            "element": self.element.get_struct(simplified),
+            "target": self.target.get_struct(simplified),
+        }
+        if isinstance(self.iterator, (LiteralList, LiteralTuple, LiteralSet)):
+            value["iterator"] = {
+                "type": self.iterator.__class__.__name__,
+                "elements": [e.get_struct(simplified) for e in self.iterator.elements]
+            }
+        else:
+            value["iterator"] = self.iterator.get_struct(simplified)
+
+        return self._prepare_struct(key, value, simplified)

src/astx/literals/collections.py

@@ -69,7 +69,6 @@ def __init__(
         )
         self.loc = loc
 
-
 @public
 @typechecked
 class LiteralDict(Literal):

src/astx/tools/transpilers/python.py

@@ -66,6 +66,31 @@ def visit(self, node: astx.AssignmentExpr) -> str:
         target_str = " = ".join(self.visit(target) for target in node.targets)
         return f"{target_str} = {self.visit(node.value)}"
 
+    @dispatch  # type: ignore[no-redef]
+    def visit(self, node: astx.AsyncForRangeLoopExpr) -> str:
+        """Handle AsyncForRangeLoopExpr nodes."""
+        if len(node.body) > 1:
+            raise ValueError(
+                "AsyncForRangeLoopExpr in Python just accept 1 node in the "
+                "body attribute."
+            )
+        start = (
+            self.visit(node.start)
+            if getattr(node, "start", None) is not None
+            else "0"
+        )
+        end = self.visit(node.end)
+        step = (
+            self.visit(node.step)
+            if getattr(node, "step", None) is not None
+            else "1"
+        )
+
+        return (
+            f"result = [{self.visit(node.body).strip()} async for "
+            f"{node.variable.name} in range({start}, {end}, {step})]"
+        )
+
     @dispatch  # type: ignore[no-redef]
     def visit(self, node: astx.AwaitExpr) -> str:
         """Handle AwaitExpr nodes."""
@@ -520,6 +545,12 @@ def visit(self, node: astx.YieldExpr) -> str:
         value = self.visit(node.value) if node.value else ""
         return f"yield {value}".strip()
 
+    @dispatch  # type: ignore[no-redef]
+    def visit(self, node: astx.YieldFromExpr) -> str:
+        """Handle YieldFromExpr nodes."""
+        value = self.visit(node.value)
+        return f"yield from {value}".strip()
+
     @dispatch  # type: ignore[no-redef]
     def visit(self, node: astx.Date) -> str:
         """Handle Date nodes."""
@@ -643,3 +674,8 @@ 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.GeneratorExpr) -> str:
+        """Handle GeneratorExr nodes."""
+        return f"({self.visit(node.element)} for {self.visit(node.target)} in {self.visit(node.iterator)})".strip()