fix: consider union template types on builders (#1701)

larastan · Jul 29, 2023 · 6c5e882 · 6c5e882
1 parent 5863603
commit 6c5e882
Show file tree

Hide file tree

Showing 3 changed files with 81 additions and 23 deletions.
diff --git a/phpstan-baseline.neon b/phpstan-baseline.neon
@@ -30,11 +30,6 @@ parameters:
 			count: 1
 			path: src/Properties/ModelAccessorExtension.php
 
-		-
-			message: "#^PHPDoc tag @var assumes the expression with type PHPStan\\\\Type\\\\Type\\|null is always PHPStan\\\\Type\\\\ObjectType but it's error\\-prone and dangerous\\.$#"
-			count: 1
-			path: src/ReturnTypes/BuilderModelFindExtension.php
-
 		-
 			message: "#^Constant LARAVEL_VERSION not found\\.$#"
 			count: 1

diff --git a/src/ReturnTypes/BuilderModelFindExtension.php b/src/ReturnTypes/BuilderModelFindExtension.php
@@ -6,6 +6,7 @@
 
 use Illuminate\Database\Eloquent\Builder;
 use Illuminate\Database\Eloquent\Collection;
+use Illuminate\Database\Eloquent\Model;
 use Illuminate\Database\Query\Builder as QueryBuilder;
 use Illuminate\Support\Str;
 use NunoMaduro\Larastan\Methods\ModelTypeHelper;
@@ -78,31 +79,45 @@ public function getTypeFromMethodCall(
             return null;
         }
 
-        /** @var ObjectType $model */
-        $model = $methodReflection->getDeclaringClass()->getActiveTemplateTypeMap()->getType('TModelClass');
+        /** @var Type $modelClassType */
+        $modelClassType = $methodReflection->getDeclaringClass()->getActiveTemplateTypeMap()->getType('TModelClass');
+
+        if ((new ObjectType(Model::class))->isSuperTypeOf($modelClassType)->no()) {
+            return null;
+        }
+
         $returnType = $methodReflection->getVariants()[0]->getReturnType();
         $argType = $scope->getType($methodCall->getArgs()[0]->value);
 
-        $returnType = ModelTypeHelper::replaceStaticTypeWithModel($returnType, $model->getClassName());
+        if ($argType instanceof MixedType) {
+            return $returnType;
+        }
 
-        if ($argType->isIterable()->yes()) {
-            if (in_array(Collection::class, $returnType->getReferencedClasses(), true)) {
-                return $this->collectionHelper->determineCollectionClass($model->getClassName());
-            }
+        $models = [];
 
-            return TypeCombinator::remove($returnType, $model);
-        }
+        foreach ($modelClassType->getObjectClassReflections() as $objectClassReflection) {
+            $modelName = $objectClassReflection->getName();
 
-        if ($argType instanceof MixedType) {
-            return $returnType;
+            $returnType = ModelTypeHelper::replaceStaticTypeWithModel($returnType, $modelName);
+
+            if ($argType->isIterable()->yes()) {
+                if (in_array(Collection::class, $returnType->getReferencedClasses(), true)) {
+                    $models[] = $this->collectionHelper->determineCollectionClass($modelName);
+                    continue;
+                }
+
+                $models[] = TypeCombinator::remove($returnType, new ObjectType($modelName));
+            } else {
+                $models[] = TypeCombinator::remove(
+                    TypeCombinator::remove(
+                        $returnType,
+                        new ArrayType(new MixedType(), $modelClassType)
+                    ),
+                    new ObjectType(Collection::class)
+                );
+            }
         }
 
-        return TypeCombinator::remove(
-            TypeCombinator::remove(
-                $returnType,
-                new ArrayType(new MixedType(), $model)
-            ),
-            new ObjectType(Collection::class)
-        );
+        return count($models) > 0 ? TypeCombinator::union(...$models) : null;
     }
 }
diff --git a/tests/Type/data/eloquent-builder.php b/tests/Type/data/eloquent-builder.php
@@ -6,6 +6,7 @@
 
 use App\Post;
 use App\PostBuilder;
+use App\Team;
 use App\User;
 use Illuminate\Database\Eloquent\Builder;
 use Illuminate\Database\Eloquent\Builder as EloquentBuilder;
@@ -451,6 +452,13 @@ function testChunkOnEloquentBuilder()
     User::chunk(1000, fn ($collection) => assertType('Illuminate\Database\Eloquent\Collection<int, App\User>', $collection));
 }
 
+/** @param Builder<User|Team> $builder */
+function testUnionBuilder(Builder $builder)
+{
+    assertType('App\Team|App\User', $builder->findOrFail(4));
+    assertType('Illuminate\Database\Eloquent\Builder<App\Team|App\User>', $builder->where('id', 5));
+}
+
 class Foo extends Model
 {
     /** @phpstan-use FooTrait<Foo> */
@@ -492,3 +500,43 @@ public function testQuery(): Builder
 class CustomBuilder extends Builder
 {
 }
+
+/**
+ * @template TModel of User|Team
+ */
+abstract class UnionClass
+{
+    /**
+     * @return TModel
+     */
+    public function test(int $id): Model
+    {
+        assertType('TModel of App\Team|App\User (class EloquentBuilder\UnionClass, argument)', $this->getQuery()->findOrFail($id));
+
+        return $this->getQuery()->findOrFail($id);
+    }
+
+    /**
+     * @return Builder<TModel>
+     */
+    abstract public function getQuery(): Builder;
+}
+
+/**
+ * @extends UnionClass<Team>
+ */
+class TeamClass extends UnionClass
+{
+    public function foo()
+    {
+        assertType('App\Team', $this->test(5));
+    }
+
+    /**
+     * {@inheritdoc}
+     */
+    public function getQuery(): Builder
+    {
+        return Team::query();
+    }
+}