Check iterator types on yield from (#5042)

* Check iterator types on yield from

* Switch to NodeAbstract

* Make Foo iterable

src/Psalm/Internal/Analyzer/Statements/Block/ForeachAnalyzer.php

@@ -215,6 +215,7 @@ class ForeachAnalyzer
             if (self::checkIteratorType(
                 $statements_analyzer,
                 $stmt,
+                $stmt->expr,
                 $iterator_type,
                 $codebase,
                 $context,
@@ -335,11 +336,13 @@ class ForeachAnalyzer
     }
 
     /**
+     * @param PhpParser\Node\Stmt\Foreach_|PhpParser\Node\Expr\YieldFrom $stmt
      * @return false|null
      */
     public static function checkIteratorType(
         StatementsAnalyzer $statements_analyzer,
-        PhpParser\Node\Stmt\Foreach_ $stmt,
+        PhpParser\NodeAbstract $stmt,
+        PhpParser\Node\Expr $expr,
         Type\Union $iterator_type,
         Codebase $codebase,
         Context $context,
@@ -351,7 +354,7 @@ class ForeachAnalyzer
             if (IssueBuffer::accepts(
                 new NullIterator(
                     'Cannot iterate over null',
-                    new CodeLocation($statements_analyzer->getSource(), $stmt->expr)
+                    new CodeLocation($statements_analyzer->getSource(), $expr)
                 ),
                 $statements_analyzer->getSuppressedIssues()
             )) {
@@ -361,7 +364,7 @@ class ForeachAnalyzer
             if (IssueBuffer::accepts(
                 new PossiblyNullIterator(
                     'Cannot iterate over nullable var ' . $iterator_type,
-                    new CodeLocation($statements_analyzer->getSource(), $stmt->expr)
+                    new CodeLocation($statements_analyzer->getSource(), $expr)
                 ),
                 $statements_analyzer->getSuppressedIssues()
             )) {
@@ -371,7 +374,7 @@ class ForeachAnalyzer
             if (IssueBuffer::accepts(
                 new PossiblyFalseIterator(
                     'Cannot iterate over falsable var ' . $iterator_type,
-                    new CodeLocation($statements_analyzer->getSource(), $stmt->expr)
+                    new CodeLocation($statements_analyzer->getSource(), $expr)
                 ),
                 $statements_analyzer->getSuppressedIssues()
             )) {
@@ -415,7 +418,7 @@ class ForeachAnalyzer
                     $iterator_atomic_type = $iterator_atomic_type->getGenericArrayType();
                 } elseif ($iterator_atomic_type instanceof Type\Atomic\TList) {
                     $list_var_id = ExpressionIdentifier::getArrayVarId(
-                        $stmt->expr,
+                        $expr,
                         $statements_analyzer->getFQCLN(),
                         $statements_analyzer
                     );
@@ -444,7 +447,7 @@ class ForeachAnalyzer
 
                 ArrayFetchAnalyzer::taintArrayFetch(
                     $statements_analyzer,
-                    $stmt->expr,
+                    $expr,
                     null,
                     $value_type,
                     Type::getMixed()
@@ -479,7 +482,7 @@ class ForeachAnalyzer
 
                 ArrayFetchAnalyzer::taintArrayFetch(
                     $statements_analyzer,
-                    $stmt->expr,
+                    $expr,
                     null,
                     $value_type,
                     Type::getMixed()
@@ -592,7 +595,7 @@ class ForeachAnalyzer
                     if (ClassLikeAnalyzer::checkFullyQualifiedClassLikeName(
                         $statements_analyzer,
                         $iterator_atomic_type->value,
-                        new CodeLocation($statements_analyzer->getSource(), $stmt->expr),
+                        new CodeLocation($statements_analyzer->getSource(), $expr),
                         $context->self,
                         $context->calling_method_id,
                         $statements_analyzer->getSuppressedIssues()
@@ -609,7 +612,7 @@ class ForeachAnalyzer
                     self::handleIterable(
                         $statements_analyzer,
                         $iterator_atomic_type,
-                        $stmt->expr,
+                        $expr,
                         $codebase,
                         $context,
                         $key_type,
@@ -647,7 +650,7 @@ class ForeachAnalyzer
                 if (IssueBuffer::accepts(
                     new PossibleRawObjectIteration(
                         'Possibly undesired iteration over regular object ' . \reset($raw_object_types),
-                        new CodeLocation($statements_analyzer->getSource(), $stmt->expr)
+                        new CodeLocation($statements_analyzer->getSource(), $expr)
                     ),
                     $statements_analyzer->getSuppressedIssues()
                 )) {
@@ -657,7 +660,7 @@ class ForeachAnalyzer
                 if (IssueBuffer::accepts(
                     new RawObjectIteration(
                         'Possibly undesired iteration over regular object ' . \reset($raw_object_types),
-                        new CodeLocation($statements_analyzer->getSource(), $stmt->expr)
+                        new CodeLocation($statements_analyzer->getSource(), $expr)
                     ),
                     $statements_analyzer->getSuppressedIssues()
                 )) {
@@ -671,7 +674,7 @@ class ForeachAnalyzer
                 if (IssueBuffer::accepts(
                     new PossiblyInvalidIterator(
                         'Cannot iterate over ' . $invalid_iterator_types[0],
-                        new CodeLocation($statements_analyzer->getSource(), $stmt->expr)
+                        new CodeLocation($statements_analyzer->getSource(), $expr)
                     ),
                     $statements_analyzer->getSuppressedIssues()
                 )) {
@@ -681,7 +684,7 @@ class ForeachAnalyzer
                 if (IssueBuffer::accepts(
                     new InvalidIterator(
                         'Cannot iterate over ' . $invalid_iterator_types[0],
-                        new CodeLocation($statements_analyzer->getSource(), $stmt->expr)
+                        new CodeLocation($statements_analyzer->getSource(), $expr)
                     ),
                     $statements_analyzer->getSuppressedIssues()
                 )) {

src/Psalm/Internal/Analyzer/Statements/Expression/YieldFromAnalyzer.php

@@ -5,6 +5,7 @@ use PhpParser;
 use Psalm\Internal\Analyzer\Statements\ExpressionAnalyzer;
 use Psalm\Internal\Analyzer\StatementsAnalyzer;
 use Psalm\Context;
+use Psalm\Internal\Analyzer\Statements\Block\ForeachAnalyzer;
 use Psalm\Type;
 use function strtolower;
 
@@ -26,6 +27,24 @@ class YieldFromAnalyzer
         }
 
         if ($stmt_expr_type = $statements_analyzer->node_data->getType($stmt->expr)) {
+            $key_type = null;
+            $value_type = null;
+            $always_non_empty_array = true;
+            if (ForeachAnalyzer::checkIteratorType(
+                $statements_analyzer,
+                $stmt,
+                $stmt->expr,
+                $stmt_expr_type,
+                $statements_analyzer->getCodebase(),
+                $context,
+                $key_type,
+                $value_type,
+                $always_non_empty_array
+            ) === false
+            ) {
+                return false;
+            }
+            
             $yield_from_type = null;
 
             foreach ($stmt_expr_type->getAtomicTypes() as $atomic_type) {

tests/GeneratorTest.php

@@ -314,6 +314,63 @@ class GeneratorTest extends TestCase
                     }',
                 'error_message' => 'NullableReturnStatement',
             ],
+            'invalidIterator' => [
+                '<?php
+                    function example() : int {
+                        return 0;
+                    }
+
+                    function example2() : Generator {
+                        yield from example();
+                    }',
+                'error_message' => 'InvalidIterator',
+            ],
+            'rawObjectIteration' => [
+                '<?php
+                    class A {
+                        /** @var ?string */
+                        public $foo;
+                    }
+                    function example() : Generator {
+                        $arr = new A;
+
+                        yield from $arr;
+                    }',
+                'error_message' => 'RawObjectIteration',
+            ],
+            'possibleRawObjectIteration' => [
+                '<?php
+                    class A {
+                        /** @var ?string */
+                        public $foo;
+                    }
+
+                    class B extends A {}
+
+                    function bar(A $a): void {}
+
+                    function gen() : Generator {
+                        $arr = [];
+
+                        if (rand(0, 10) > 5) {
+                            $arr[] = new A;
+                        } else {
+                            $arr = new B;
+                        }
+
+                        yield from $arr;
+                    }',
+                'error_message' => 'PossibleRawObjectIteration',
+            ],
+            'possibleRawObjectIterationFromIsset' => [
+                '<?php
+                    function foo(array $a) : Generator {
+                        if (isset($a["a"]["b"])) {
+                            yield from $a["a"];
+                        }
+                    }',
+                'error_message' => 'PossibleRawObjectIteration',
+            ],
         ];
     }
 }

tests/Template/ClassTemplateTest.php

@@ -2024,7 +2024,8 @@ class ClassTemplateTest extends TestCase
             ],
             'yieldFromGenericObjectNotExtendingIterator' => [
                 '<?php
-                    class Foo{}
+                    /** @extends \ArrayObject<int, int> */
+                    class Foo extends \ArrayObject {}
 
                     class A {
                         /**