From 27f6171f2f98e7b9c0d9e175118b23177a02d414 Mon Sep 17 00:00:00 2001
From: Fabien Potencier <fabien@potencier.org>
Date: Mon, 22 Jul 2024 13:46:32 +0200
Subject: [PATCH] Add support for yielding from a generator in PrintNode

---
 CHANGELOG                                  |  1 +
 src/Node/Expression/AbstractExpression.php |  4 +++
 src/Node/PrintNode.php                     |  5 ++-
 src/NodeVisitor/SandboxNodeVisitor.php     |  2 +-
 tests/Node/PrintTest.php                   |  9 +++++
 tests/NodeVisitor/SandboxTest.php          | 40 ++++++++++++++++++++++
 6 files changed, 57 insertions(+), 4 deletions(-)
 create mode 100644 tests/NodeVisitor/SandboxTest.php

diff --git a/CHANGELOG b/CHANGELOG
index 933fec8fb..785a26949 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -1,5 +1,6 @@
 # 3.11.0 (2024-XX-XX)
 
+ * Add the possibility to yield from a generator in `PrintNode`
  * Add the `shuffle` filter
  * Add the `singular` and `plural` filters in `StringExtension`
  * Deprecate the second argument of `Twig\Node\Expression\CallExpression::compileArguments()`
diff --git a/src/Node/Expression/AbstractExpression.php b/src/Node/Expression/AbstractExpression.php
index 42da0559d..1692f5671 100644
--- a/src/Node/Expression/AbstractExpression.php
+++ b/src/Node/Expression/AbstractExpression.php
@@ -21,4 +21,8 @@ use Twig\Node\Node;
  */
 abstract class AbstractExpression extends Node
 {
+    public function isGenerator(): bool
+    {
+        return $this->hasAttribute('is_generator') && $this->getAttribute('is_generator');
+    }
 }
diff --git a/src/Node/PrintNode.php b/src/Node/PrintNode.php
index a6a89bd74..bdc738301 100644
--- a/src/Node/PrintNode.php
+++ b/src/Node/PrintNode.php
@@ -31,10 +31,9 @@ class PrintNode extends Node implements NodeOutputInterface
 
     public function compile(Compiler $compiler): void
     {
-        $compiler->addDebugInfo($this);
-
         $compiler
-            ->write('yield ')
+            ->addDebugInfo($this)
+            ->write($this->getNode('expr')->isGenerator() ? 'yield from ' : 'yield ')
             ->subcompile($this->getNode('expr'))
             ->raw(";\n")
         ;
diff --git a/src/NodeVisitor/SandboxNodeVisitor.php b/src/NodeVisitor/SandboxNodeVisitor.php
index d1108394f..68020885e 100644
--- a/src/NodeVisitor/SandboxNodeVisitor.php
+++ b/src/NodeVisitor/SandboxNodeVisitor.php
@@ -119,7 +119,7 @@ final class SandboxNodeVisitor implements NodeVisitorInterface
     private function wrapNode(Node $node, string $name): void
     {
         $expr = $node->getNode($name);
-        if ($expr instanceof NameExpression || $expr instanceof GetAttrExpression) {
+        if (($expr instanceof NameExpression || $expr instanceof GetAttrExpression) && !$expr->isGenerator()) {
             $node->setNode($name, new CheckToStringNode($expr));
         }
     }
diff --git a/tests/Node/PrintTest.php b/tests/Node/PrintTest.php
index 2df440c28..09c2a19ab 100644
--- a/tests/Node/PrintTest.php
+++ b/tests/Node/PrintTest.php
@@ -12,7 +12,10 @@ namespace Twig\Tests\Node;
  */
 
 use Twig\Node\Expression\ConstantExpression;
+use Twig\Node\Expression\GetAttrExpression;
+use Twig\Node\Expression\NameExpression;
 use Twig\Node\PrintNode;
+use Twig\Template;
 use Twig\Test\NodeTestCase;
 
 class PrintTest extends NodeTestCase
@@ -30,6 +33,12 @@ class PrintTest extends NodeTestCase
         $tests = [];
         $tests[] = [new PrintNode(new ConstantExpression('foo', 1), 1), "// line 1\nyield \"foo\";"];
 
+        $expr = new NameExpression('foo', 1);
+        $attr = new ConstantExpression('bar', 1);
+        $node = new GetAttrExpression($expr, $attr, null, Template::METHOD_CALL, 1);
+        $node->setAttribute('is_generator', true);
+        $tests[] = [new PrintNode($node, 1), "// line 1\nyield from CoreExtension::getAttribute(\$this->env, \$this->source, (\$context[\"foo\"] ?? null), \"bar\", [], \"method\", false, false, false, 1);"];
+
         return $tests;
     }
 }
diff --git a/tests/NodeVisitor/SandboxTest.php b/tests/NodeVisitor/SandboxTest.php
new file mode 100644
index 000000000..d465abc35
--- /dev/null
+++ b/tests/NodeVisitor/SandboxTest.php
@@ -0,0 +1,40 @@
+<?php
+
+namespace Twig\Tests\NodeVisitor;
+
+/*
+ * This file is part of Twig.
+ *
+ * (c) Fabien Potencier
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+use PHPUnit\Framework\TestCase;
+use Twig\Environment;
+use Twig\Loader\ArrayLoader;
+use Twig\Node\CheckToStringNode;
+use Twig\Node\Expression\NameExpression;
+use Twig\Node\ModuleNode;
+use Twig\Node\Node;
+use Twig\Node\PrintNode;
+use Twig\NodeTraverser;
+use Twig\NodeVisitor\SandboxNodeVisitor;
+use Twig\Source;
+
+class SandboxTest extends TestCase
+{
+    public function testGeneratorExpression()
+    {
+        $env = new Environment(new ArrayLoader());
+        $expr = new NameExpression('foo', 1);
+        $expr->setAttribute('is_generator', true);
+        $node = new ModuleNode(new PrintNode($expr, 1), null, new Node(), new Node(), new Node(), new Node([]), new Source('foo', 'foo'));
+        $traverser = new NodeTraverser($env, [new SandboxNodeVisitor($env)]);
+        $node = $traverser->traverse($node);
+
+        $this->assertNotInstanceOf(CheckToStringNode::class, $node->getNode('body')->getNode('expr'));
+        $this->assertSame("// line 1\nyield from (\$context[\"foo\"] ?? null);\n", $env->compile($node->getNode('body')));
+    }
+}