fix: resolve $this type in ->with(), dataset closures, and trait-bound test cases

extension.neon

@@ -41,11 +41,24 @@ services:
         tags:
             - phpstan.broker.dynamicMethodReturnTypeExtension
 
+    -
+        class: Pest\PHPStan\Type\Pest\PestTestCaseType
+
     -
         class: Pest\PHPStan\Type\Pest\TestClosureThisTypeExtension
         tags:
             - phpstan.functionParameterClosureThisExtension
 
+    -
+        class: Pest\PHPStan\Type\Pest\WithClosureThisTypeExtension
+        tags:
+            - phpstan.methodParameterClosureThisExtension
+
+    -
+        class: Pest\PHPStan\Type\Pest\WithDatasetClosureNodeVisitor
+        tags:
+            - phpstan.parser.richParserNodeVisitor
+
     -
         class: Pest\PHPStan\Type\Pest\ExpectationPropertiesExtension
         tags:

src/Type/Pest/PestConfigReader.php

@@ -24,6 +24,9 @@ final class PestConfigReader
     /** @var array<string, list<array{class: string, config: string}>> */
     private array $globalUseDirectoryMap = [];
 
+    /** @var list<string> Bindings declared in a Pest.php config without an ->in() scope */
+    private array $globalBindings = [];
+
     /** @var array<string, list<string>> Caches bindings declared directly inside a test file */
     private array $fileBindingsCache = [];
 
@@ -45,7 +48,7 @@ public function resolveBindings(string $filePath): array
         $this->ensureParsed();
 
         $normalizedFile = $this->fileDiscoverer->normalizePath($filePath);
-        $bindings = [];
+        $bindings = [...$this->globalBindings];
 
         foreach ($this->directoryMap as $bindingKey => $classNames) {
             if (! $this->targetResolver->matches($bindingKey, $normalizedFile)) {
@@ -120,7 +123,7 @@ public function allBoundClasses(): array
     {
         $this->ensureParsed();
 
-        $bindings = [];
+        $bindings = [...$this->globalBindings];
 
         foreach ($this->directoryMap as $classNames) {
             array_push($bindings, ...$classNames);
@@ -158,6 +161,27 @@ private function parsePestFile(string $filePath): void
 
         $this->extractUsesBindings($nodeFinder, $stmts, $pestFileDir, $filePath);
         $this->extractPestBindings($nodeFinder, $stmts, $pestFileDir, $filePath);
+        $this->extractGlobalBindings($stmts);
+    }
+
+    /**
+     * @param  Node[]  $stmts
+     */
+    private function extractGlobalBindings(array $stmts): void
+    {
+        foreach ($this->topLevelExpressions($stmts) as $expr) {
+            $usesArgs = $this->extractFileUsesArgs($expr);
+            if ($usesArgs !== []) {
+                array_push($this->globalBindings, ...$usesArgs);
+
+                continue;
+            }
+
+            $pestArgs = $this->extractFilePestArgs($expr);
+            if ($pestArgs !== []) {
+                array_push($this->globalBindings, ...$pestArgs);
+            }
+        }
     }
 
     /**

src/Type/Pest/PestTestCaseType.php

@@ -0,0 +1,78 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Pest\PHPStan\Type\Pest;
+
+use PHPStan\Reflection\ReflectionProvider;
+use PHPStan\Type\ObjectType;
+use PHPStan\Type\Type;
+use PHPStan\Type\TypeCombinator;
+use PHPUnit\Framework\TestCase;
+
+final class PestTestCaseType
+{
+    public function __construct(
+        private readonly PestConfigReader $pestConfigReader,
+        private readonly ReflectionProvider $reflectionProvider,
+    ) {}
+
+    public function resolve(string $filePath): Type
+    {
+        $bindings = $this->pestConfigReader->resolveFileBindings($filePath);
+
+        if ($bindings === []) {
+            $bindings = $this->pestConfigReader->resolveBindings($filePath);
+        }
+
+        $classNames = [];
+        $traitNames = [];
+
+        foreach ($bindings as $binding) {
+            if (! $this->reflectionProvider->hasClass($binding)) {
+                continue;
+            }
+
+            $reflection = $this->reflectionProvider->getClass($binding);
+
+            if ($reflection->isTrait()) {
+                $traitNames[] = $binding;
+
+                continue;
+            }
+
+            $classNames[] = $binding;
+        }
+
+        if ($classNames === []) {
+            $classNames[] = TestCase::class;
+        }
+
+        $classType = $this->toObjectType($classNames);
+
+        if ($traitNames === []) {
+            return $classType;
+        }
+
+        return new PestTestCaseWithTraitsType(
+            $classNames[0],
+            $traitNames,
+            $this->reflectionProvider,
+        );
+    }
+
+    /**
+     * @param  list<string>  $classNames
+     */
+    private function toObjectType(array $classNames): Type
+    {
+        if (count($classNames) === 1) {
+            return new ObjectType($classNames[0]);
+        }
+
+        return TypeCombinator::intersect(...array_map(
+            static fn (string $className): ObjectType => new ObjectType($className),
+            $classNames,
+        ));
+    }
+}

src/Type/Pest/PestTestCaseWithTraitsType.php

@@ -0,0 +1,171 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Pest\PHPStan\Type\Pest;
+
+use Override;
+use PHPStan\Reflection\ClassConstantReflection;
+use PHPStan\Reflection\ClassMemberAccessAnswerer;
+use PHPStan\Reflection\ExtendedMethodReflection;
+use PHPStan\Reflection\ExtendedPropertyReflection;
+use PHPStan\Reflection\ReflectionProvider;
+use PHPStan\Reflection\Type\UnresolvedMethodPrototypeReflection;
+use PHPStan\Reflection\Type\UnresolvedPropertyPrototypeReflection;
+use PHPStan\TrinaryLogic;
+use PHPStan\Type\ObjectType;
+
+final class PestTestCaseWithTraitsType extends ObjectType
+{
+    /** @param  list<class-string>  $traitNames */
+    public function __construct(
+        string $className,
+        private readonly array $traitNames,
+        private readonly ReflectionProvider $reflectionProvider,
+    ) {
+        parent::__construct($className);
+    }
+
+    #[Override]
+    public function hasMethod(string $methodName): TrinaryLogic
+    {
+        if (parent::hasMethod($methodName)->yes()) {
+            return TrinaryLogic::createYes();
+        }
+
+        if ($this->hasTraitMethod($methodName)) {
+            return TrinaryLogic::createYes();
+        }
+
+        return parent::hasMethod($methodName);
+    }
+
+    #[Override]
+    public function getMethod(string $methodName, ClassMemberAccessAnswerer $scope): ExtendedMethodReflection
+    {
+        if (parent::hasMethod($methodName)->yes()) {
+            return parent::getMethod($methodName, $scope);
+        }
+
+        foreach ($this->traitNames as $traitName) {
+            $traitReflection = $this->reflectionProvider->getClass($traitName);
+
+            if ($traitReflection->hasNativeMethod($methodName)) {
+                return $traitReflection->getNativeMethod($methodName);
+            }
+        }
+
+        return parent::getMethod($methodName, $scope);
+    }
+
+    #[Override]
+    public function getUnresolvedMethodPrototype(string $methodName, ClassMemberAccessAnswerer $scope): UnresolvedMethodPrototypeReflection
+    {
+        if (parent::hasMethod($methodName)->yes()) {
+            return parent::getUnresolvedMethodPrototype($methodName, $scope);
+        }
+
+        foreach ($this->traitNames as $traitName) {
+            if ($this->reflectionProvider->getClass($traitName)->hasNativeMethod($methodName)) {
+                return new ObjectType($traitName)->getUnresolvedMethodPrototype($methodName, $scope);
+            }
+        }
+
+        return parent::getUnresolvedMethodPrototype($methodName, $scope);
+    }
+
+    #[Override]
+    public function hasProperty(string $propertyName): TrinaryLogic
+    {
+        if (parent::hasProperty($propertyName)->yes()) {
+            return TrinaryLogic::createYes();
+        }
+
+        if ($this->hasTraitProperty($propertyName)) {
+            return TrinaryLogic::createYes();
+        }
+
+        return parent::hasProperty($propertyName);
+    }
+
+    #[Override]
+    public function getProperty(string $propertyName, ClassMemberAccessAnswerer $scope): ExtendedPropertyReflection
+    {
+        if (parent::hasProperty($propertyName)->yes()) {
+            return parent::getProperty($propertyName, $scope);
+        }
+
+        foreach ($this->traitNames as $traitName) {
+            $traitReflection = $this->reflectionProvider->getClass($traitName);
+
+            if ($traitReflection->hasNativeProperty($propertyName)) {
+                return $traitReflection->getNativeProperty($propertyName);
+            }
+        }
+
+        return parent::getProperty($propertyName, $scope);
+    }
+
+    #[Override]
+    public function getUnresolvedPropertyPrototype(string $propertyName, ClassMemberAccessAnswerer $scope): UnresolvedPropertyPrototypeReflection
+    {
+        if (parent::hasProperty($propertyName)->yes()) {
+            return parent::getUnresolvedPropertyPrototype($propertyName, $scope);
+        }
+
+        foreach ($this->traitNames as $traitName) {
+            if ($this->reflectionProvider->getClass($traitName)->hasNativeProperty($propertyName)) {
+                return new ObjectType($traitName)->getUnresolvedPropertyPrototype($propertyName, $scope);
+            }
+        }
+
+        return parent::getUnresolvedPropertyPrototype($propertyName, $scope);
+    }
+
+    #[Override]
+    public function hasConstant(string $constantName): TrinaryLogic
+    {
+        if (parent::hasConstant($constantName)->yes()) {
+            return TrinaryLogic::createYes();
+        }
+
+        if ($this->hasTraitConstant($constantName)) {
+            return TrinaryLogic::createYes();
+        }
+
+        return parent::hasConstant($constantName);
+    }
+
+    #[Override]
+    public function getConstant(string $constantName): ClassConstantReflection
+    {
+        if (parent::hasConstant($constantName)->yes()) {
+            return parent::getConstant($constantName);
+        }
+
+        foreach ($this->traitNames as $traitName) {
+            $traitReflection = $this->reflectionProvider->getClass($traitName);
+
+            if ($traitReflection->hasConstant($constantName)) {
+                return $traitReflection->getConstant($constantName);
+            }
+        }
+
+        return parent::getConstant($constantName);
+    }
+
+    private function hasTraitMethod(string $methodName): bool
+    {
+        return array_any($this->traitNames, fn (string $traitName): bool => $this->reflectionProvider->getClass($traitName)->hasNativeMethod($methodName));
+    }
+
+    private function hasTraitProperty(string $propertyName): bool
+    {
+        return array_any($this->traitNames, fn (string $traitName): bool => $this->reflectionProvider->getClass($traitName)->hasNativeProperty($propertyName));
+    }
+
+    private function hasTraitConstant(string $constantName): bool
+    {
+        return array_any($this->traitNames, fn (string $traitName): bool => $this->reflectionProvider->getClass($traitName)->hasConstant($constantName));
+    }
+}

src/Type/Pest/TestClosureThisTypeExtension.php

@@ -8,12 +8,8 @@
 use PHPStan\Analyser\Scope;
 use PHPStan\Reflection\FunctionReflection;
 use PHPStan\Reflection\ParameterReflection;
-use PHPStan\Reflection\ReflectionProvider;
 use PHPStan\Type\FunctionParameterClosureThisExtension;
-use PHPStan\Type\ObjectType;
 use PHPStan\Type\Type;
-use PHPStan\Type\TypeCombinator;
-use PHPUnit\Framework\TestCase;
 
 final class TestClosureThisTypeExtension implements FunctionParameterClosureThisExtension
 {
@@ -31,8 +27,7 @@ final class TestClosureThisTypeExtension implements FunctionParameterClosureThis
     ];
 
     public function __construct(
-        private readonly PestConfigReader $pestConfigReader,
-        private readonly ReflectionProvider $reflectionProvider,
+        private readonly PestTestCaseType $pestTestCaseType,
     ) {}
 
     public function isFunctionSupported(FunctionReflection $functionReflection, ParameterReflection $parameter): bool
@@ -49,43 +44,6 @@ public function getClosureThisTypeFromFunctionCall(
         ParameterReflection $parameter,
         Scope $scope
     ): Type {
-        $types = $this->toClassObjectTypes(
-            $this->pestConfigReader->resolveFileBindings($scope->getFile()),
-        );
-
-        if ($types === []) {
-            $types = $this->toClassObjectTypes(
-                $this->pestConfigReader->resolveBindings($scope->getFile()),
-            );
-        }
-
-        if ($types === []) {
-            return new ObjectType(TestCase::class);
-        }
-
-        return count($types) === 1 ? $types[0] : TypeCombinator::intersect(...$types);
-    }
-
-    /**
-     * @param  list<string>  $bindings
-     * @return list<ObjectType>
-     */
-    private function toClassObjectTypes(array $bindings): array
-    {
-        $types = [];
-
-        foreach ($bindings as $binding) {
-            if (! $this->reflectionProvider->hasClass($binding)) {
-                continue;
-            }
-
-            if ($this->reflectionProvider->getClass($binding)->isTrait()) {
-                continue;
-            }
-
-            $types[] = new ObjectType($binding);
-        }
-
-        return $types;
+        return $this->pestTestCaseType->resolve($scope->getFile());
     }
 }