Handle more iterable types

klimick · May 3, 2024 · b423be7 · b423be7
1 parent 223d07f
commit b423be7
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/src/Psalm/Internal/Type/TemplateContextualBoundsCollector.php b/src/Psalm/Internal/Type/TemplateContextualBoundsCollector.php
@@ -6,6 +6,7 @@
 
 use Psalm\Codebase;
 use Psalm\Internal\Type\Comparator\CallableTypeComparator;
+use Psalm\Type;
 use Psalm\Type\Atomic;
 use Psalm\Type\Atomic\TArray;
 use Psalm\Type\Atomic\TCallable;
@@ -162,7 +163,7 @@ private function collectKeyedArray(TKeyedArray $contextual_atomic, Atomic $retur
                 $this->collectUnion($contextual_atomic->fallback_params[0], $return_atomic->fallback_params[0]);
                 $this->collectUnion($contextual_atomic->fallback_params[1], $return_atomic->fallback_params[1]);
             }
-        } elseif ($return_atomic instanceof TArray) {
+        } elseif ($return_atomic instanceof TArray || $return_atomic instanceof TIterable) {
             $this->collectUnion($contextual_atomic->getGenericKeyType(), $return_atomic->type_params[0]);
             $this->collectUnion($contextual_atomic->getGenericValueType(), $return_atomic->type_params[1]);
         }
@@ -200,6 +201,15 @@ private function collectGenericObject(TGenericObject $contextual_atomic, Atomic
             );
         }
 
+        if ($return_atomic instanceof TIterable
+            && $contextual_atomic->value === 'Generator'
+        ) {
+            $this->collectAtomic(new TIterable([
+                $contextual_atomic->type_params[0] ?? Type::getMixed(),
+                $contextual_atomic->type_params[1] ?? Type::getMixed(),
+            ]), $return_atomic);
+        }
+
         if ($return_atomic instanceof TGenericObject
             && $contextual_atomic->value === $return_atomic->value
         ) {