Leverage custom object type for county entity queries

Author: mglaman

extension.neon

@@ -44,7 +44,7 @@ services:
 		class: mglaman\PHPStanDrupal\Type\DrupalClassResolverDynamicReturnTypeExtension
 		tags: [phpstan.broker.dynamicMethodReturnTypeExtension]
 	-
-		class: mglaman\PHPStanDrupal\Type\EntityQueryDynamicReturnTypeExtension
+		class: mglaman\PHPStanDrupal\Type\EntityQuery\EntityQueryDynamicReturnTypeExtension
 		tags: [phpstan.broker.dynamicMethodReturnTypeExtension]
 	-
 		class: mglaman\PHPStanDrupal\Type\DrupalClassResolverDynamicStaticReturnTypeExtension

src/Type/EntityQuery/EntityQueryCountType.php

@@ -0,0 +1,13 @@
+<?php declare(strict_types=1);
+
+namespace mglaman\PHPStanDrupal\Type\EntityQuery;
+
+use PHPStan\Type\ObjectType;
+
+/**
+ * Type used to represent an entity query instance as count query.
+ */
+final class EntityQueryCountType extends ObjectType
+{
+
+}

src/Type/EntityQueryDynamicReturnTypeExtension.php renamed to src/Type/EntityQuery/EntityQueryDynamicReturnTypeExtension.php

@@ -1,10 +1,9 @@
 <?php declare(strict_types=1);
 
-namespace mglaman\PHPStanDrupal\Type;
+namespace mglaman\PHPStanDrupal\Type\EntityQuery;
 
 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;
@@ -25,7 +24,10 @@ public function getClass(): string
 
     public function isMethodSupported(MethodReflection $methodReflection): bool
     {
-        return $methodReflection->getName() === 'execute';
+        return in_array($methodReflection->getName(), [
+            'count',
+            'execute',
+        ], true);
     }
 
     public function getTypeFromMethodCall(
@@ -34,28 +36,31 @@ public function getTypeFromMethodCall(
         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;
+        $varType = $scope->getType($methodCall->var);
+        $methodName = $methodReflection->getName();
+        if ($methodName === 'count') {
+            if ($varType instanceof ObjectType) {
+                return new EntityQueryCountType(
+                    $varType->getClassName(),
+                    $varType->getSubtractedType(),
+                    $varType->getClassReflection()
+                );
             }
+            return $defaultReturnType;
+        }
 
-            $name = $nameObject->toString();
-            if ($name === 'count') {
+        if ($methodName === 'execute') {
+            if ($varType instanceof EntityQueryCountType) {
                 return new IntegerType();
             }
-
-            $var = $var->var;
+            if ($varType instanceof ObjectType) {
+                // @todo if this is a config storage, it'd string keys.
+                // revisit after https://github.com/mglaman/phpstan-drupal/pull/239
+                // then we can check what kind of storage we have.
+                return new ArrayType(new IntegerType(), new StringType());
+            }
+            return $defaultReturnType;
         }
-
         return $defaultReturnType;
     }
 }

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

@@ -3,7 +3,11 @@
 namespace DrupalEntity;
 
 use function PHPStan\Testing\assertType;
-
+assertType(
+    'array<int, string>',
+    \Drupal::entityTypeManager()->getStorage('node')->getQuery()
+        ->execute()
+);
 assertType(
     'array<int, string>',
     \Drupal::entityTypeManager()->getStorage('node')->getQuery()
@@ -17,11 +21,23 @@
         ->count()
         ->execute()
 );
+assertType(
+    'int',
+    \Drupal::entityQuery('node')
+        ->accessCheck(TRUE)
+        ->count()
+        ->execute()
+);
+assertType(
+    'array<int, string>',
+    \Drupal::entityQuery('node')
+        ->accessCheck(TRUE)
+        ->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());
+$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());