Classes for dynamic return type extensions are found only on direct types

Author: ondrejmirtes

src/Analyser/Scope.php

@@ -1112,7 +1112,7 @@ private function resolveType(Expr $node): Type
 
 if ($node instanceof MethodCall && $node->name instanceof Node\Identifier) {
 $methodCalledOnType = $this->getType($node->var);
-$referencedClasses = $methodCalledOnType->getReferencedClasses();
+$referencedClasses = TypeUtils::getDirectClassNames($methodCalledOnType);
 $resolvedTypes = [];
 foreach ($referencedClasses as $referencedClass) {
 if (!$this->broker->hasClass($referencedClass)) {
@@ -1175,7 +1175,7 @@ private function resolveType(Expr $node): Type
 return new ErrorType();
 }
 $staticMethodReflection = $calleeType->getMethod((string) $node->name, $this);
-$referencedClasses = $calleeType->getReferencedClasses();
+$referencedClasses = TypeUtils::getDirectClassNames($calleeType);
 
 $resolvedTypes = [];
 foreach ($referencedClasses as $referencedClass) {

src/Type/TypeUtils.php

@@ -48,6 +48,32 @@ public static function generalizeType(Type $type): Type
 return $type;
 }
 
+/**
+ * @param Type $type
+ * @return string[]
+ */
+public static function getDirectClassNames(Type $type): array
+{
+if ($type instanceof TypeWithClassName) {
+return [$type->getClassName()];
+}
+
+if ($type instanceof UnionType || $type instanceof IntersectionType) {
+$classNames = [];
+foreach ($type->getTypes() as $innerType) {
+if (!$innerType instanceof TypeWithClassName) {
+continue;
+}
+
+$classNames[] = $innerType->getClassName();
+}
+
+return $classNames;
+}
+
+return [];
+}
+
 /**
  * @param Type $type
  * @return \PHPStan\Type\ConstantScalarType[]

tests/PHPStan/Analyser/NodeScopeResolverTest.php

@@ -3471,6 +3471,75 @@ public function getTypeFromStaticMethodCall(MethodReflection $methodReflection,
 );
 }
 
+public function dataDynamicReturnTypeExtensionsOnCompoundTypes(): array
+{
+return [
+[
+'DynamicMethodReturnCompoundTypes\Collection',
+'$collection->getSelf()',
+],
+[
+'DynamicMethodReturnCompoundTypes\Collection|DynamicMethodReturnCompoundTypes\Foo',
+'$collectionOrFoo->getSelf()',
+],
+];
+}
+
+/**
+ * @dataProvider dataDynamicReturnTypeExtensionsOnCompoundTypes
+ * @param string $description
+ * @param string $expression
+ */
+public function testDynamicReturnTypeExtensionsOnCompoundTypes(
+string $description,
+string $expression
+): void
+{
+$this->assertTypes(
+__DIR__ . '/data/dynamic-method-return-compound-types.php',
+$description,
+$expression,
+[
+new class () implements DynamicMethodReturnTypeExtension {
+
+public function getClass(): string
+{
+return \DynamicMethodReturnCompoundTypes\Collection::class;
+}
+
+public function isMethodSupported(MethodReflection $methodReflection): bool
+{
+return $methodReflection->getName() === 'getSelf';
+}
+
+public function getTypeFromMethodCall(MethodReflection $methodReflection, MethodCall $methodCall, Scope $scope): Type
+{
+return new ObjectType(\DynamicMethodReturnCompoundTypes\Collection::class);
+}
+
+},
+new class () implements DynamicMethodReturnTypeExtension {
+
+public function getClass(): string
+{
+return \DynamicMethodReturnCompoundTypes\Foo::class;
+}
+
+public function isMethodSupported(MethodReflection $methodReflection): bool
+{
+return $methodReflection->getName() === 'getSelf';
+}
+
+public function getTypeFromMethodCall(MethodReflection $methodReflection, MethodCall $methodCall, Scope $scope): Type
+{
+return new ObjectType(\DynamicMethodReturnCompoundTypes\Foo::class);
+}
+
+},
+]
+);
+}
+
 public function dataOverwritingVariable(): array
 {
 return [

tests/PHPStan/Analyser/data/dynamic-method-return-compound-types.php

@@ -0,0 +1,32 @@
+<?php
+
+namespace DynamicMethodReturnCompoundTypes;
+
+class Collection
+{
+
+public function getSelf()
+{
+
+}
+
+}
+
+class Foo
+{
+
+public function getSelf()
+{
+
+}
+
+/**
+ * @param Collection|Foo[] $collection
+ * @param Collection|Foo $collectionOrFoo
+ */
+public function doFoo($collection, $collectionOrFoo)
+{
+die;
+}
+
+}