Support entity query execute dynamic return type whether count() method has been called or not

Author: brambaud

extension.neon

@@ -43,6 +43,9 @@ services:
 	-
 		class: mglaman\PHPStanDrupal\Type\DrupalClassResolverDynamicReturnTypeExtension
 		tags: [phpstan.broker.dynamicMethodReturnTypeExtension]
+	-
+		class: mglaman\PHPStanDrupal\Type\EntityQueryDynamicReturnTypeExtension
+		tags: [phpstan.broker.dynamicMethodReturnTypeExtension]
 	-
 		class: mglaman\PHPStanDrupal\Type\DrupalClassResolverDynamicStaticReturnTypeExtension
 		tags: [phpstan.broker.dynamicStaticMethodReturnTypeExtension]

src/Type/EntityQueryDynamicReturnTypeExtension.php

@@ -0,0 +1,61 @@
+<?php declare(strict_types=1);
+
+namespace mglaman\PHPStanDrupal\Type;
+
+use Drupal\Core\Entity\Query\QueryInterface;
+use PhpParser\Node\Expr\MethodCall;
+use PhpParser\Node\Identifier;
+use PHPStan\Analyser\Scope;
+use PHPStan\Reflection\MethodReflection;
+use PHPStan\Reflection\ParametersAcceptorSelector;
+use PHPStan\Type\ArrayType;
+use PHPStan\Type\DynamicMethodReturnTypeExtension;
+use PHPStan\Type\IntegerType;
+use PHPStan\Type\ObjectType;
+use PHPStan\Type\StringType;
+use PHPStan\Type\Type;
+
+class EntityQueryDynamicReturnTypeExtension implements DynamicMethodReturnTypeExtension
+{
+
+    public function getClass(): string
+    {
+        return QueryInterface::class;
+    }
+
+    public function isMethodSupported(MethodReflection $methodReflection): bool
+    {
+        return $methodReflection->getName() === 'execute';
+    }
+
+    public function getTypeFromMethodCall(
+        MethodReflection $methodReflection,
+        MethodCall $methodCall,
+        Scope $scope
+    ): Type {
+        $defaultReturnType = ParametersAcceptorSelector::selectSingle($methodReflection->getVariants())->getReturnType();
+
+        $queryBuilderType = new ObjectType(QueryInterface::class);
+        $var = $methodCall->var;
+        while ($var instanceof MethodCall) {
+            $varType = $scope->getType($var->var);
+            if (!$queryBuilderType->isSuperTypeOf($varType)->yes()) {
+                return new ArrayType(new IntegerType(), new StringType());
+            }
+
+            $nameObject = $var->name;
+            if (!($nameObject instanceof Identifier)) {
+                return $defaultReturnType;
+            }
+
+            $name = $nameObject->toString();
+            if ($name === 'count') {
+                return new IntegerType();
+            }
+
+            $var = $var->var;
+        }
+
+        return $defaultReturnType;
+    }
+}

tests/src/Type/EntityQueryDynamicReturnTypeExtensionTest.php

@@ -0,0 +1,33 @@
+<?php
+
+declare(strict_types=1);
+
+namespace mglaman\PHPStanDrupal\Tests\Type;
+
+use mglaman\PHPStanDrupal\Tests\AdditionalConfigFilesTrait;
+use PHPStan\Testing\TypeInferenceTestCase;
+
+final class EntityQueryDynamicReturnTypeExtensionTest extends TypeInferenceTestCase
+{
+    use AdditionalConfigFilesTrait;
+
+    public function dataFileAsserts(): iterable
+    {
+        yield from $this->gatherAssertTypes(__DIR__ . '/data/entity-query-execute.php');
+    }
+
+    /**
+     * @dataProvider dataFileAsserts
+     * @param string $assertType
+     * @param string $file
+     * @param mixed ...$args
+     */
+    public function testFileAsserts(
+        string $assertType,
+        string $file,
+        ...$args
+    ): void
+    {
+        $this->assertFileAsserts($assertType, $file, ...$args);
+    }
+}

tests/src/Type/data/entity-query-execute.php

@@ -0,0 +1,27 @@
+<?php
+
+namespace DrupalEntity;
+
+use function PHPStan\Testing\assertType;
+
+assertType(
+    'array<int, string>',
+    \Drupal::entityTypeManager()->getStorage('node')->getQuery()
+        ->accessCheck(TRUE)
+        ->execute()
+);
+assertType(
+    'int',
+    \Drupal::entityTypeManager()->getStorage('node')->getQuery()
+        ->accessCheck(TRUE)
+        ->count()
+        ->execute()
+);
+
+// @todo: find a way support these cases.
+//$query = \Drupal::entityTypeManager()->getStorage('node')->getQuery()
+//    ->accessCheck(TRUE);
+//assertType('array<int, string>', $query->execute());
+//$query = \Drupal::entityTypeManager()->getStorage('node')->getQuery()
+//    ->accessCheck(TRUE)->count();
+//assertType('int', $query->execute());