Merge branch 'main' into typed-entity-storage-does-not-succeed-to-infer-query-type

# Conflicts:
#	src/Type/EntityStorage/EntityStorageDynamicReturnTypeExtension.php

Author: mglaman

extension.neon

@@ -297,6 +297,9 @@ services:
 	-
 		class: mglaman\PHPStanDrupal\Type\EntityStorage\EntityStorageDynamicReturnTypeExtension
 		tags: [phpstan.broker.dynamicMethodReturnTypeExtension]
+	-
+		class: mglaman\PHPStanDrupal\Type\EntityStorage\GetQueryReturnTypeExtension
+		tags: [phpstan.broker.dynamicMethodReturnTypeExtension]
 	-
 		class: mglaman\PHPStanDrupal\Type\ContainerDynamicReturnTypeExtension
 		tags: [phpstan.broker.dynamicMethodReturnTypeExtension]

src/Type/EntityStorage/EntityStorageDynamicReturnTypeExtension.php

@@ -2,23 +2,17 @@
 
 namespace mglaman\PHPStanDrupal\Type\EntityStorage;
 
-use Drupal\Core\Config\Entity\ConfigEntityStorageInterface;
-use Drupal\Core\Entity\ContentEntityStorageInterface;
 use Drupal\Core\Entity\EntityStorageInterface;
 use mglaman\PHPStanDrupal\Drupal\EntityDataRepository;
-use mglaman\PHPStanDrupal\Type\EntityQuery\ConfigEntityQueryType;
-use mglaman\PHPStanDrupal\Type\EntityQuery\ContentEntityQueryType;
 use PhpParser\Node\Expr\MethodCall;
 use PHPStan\Analyser\Scope;
 use PHPStan\Reflection\MethodReflection;
 use PHPStan\Reflection\ParametersAcceptorSelector;
-use PHPStan\ShouldNotHappenException;
 use PHPStan\Type\ArrayType;
 use PHPStan\Type\DynamicMethodReturnTypeExtension;
 use PHPStan\Type\IntegerType;
 use PHPStan\Type\ObjectType;
 use PHPStan\Type\StringType;
-use PHPStan\Type\Type;
 use PHPStan\Type\TypeCombinator;
 
 class EntityStorageDynamicReturnTypeExtension implements DynamicMethodReturnTypeExtension
@@ -49,7 +43,6 @@ public function isMethodSupported(MethodReflection $methodReflection): bool
                 'loadMultiple',
                 'loadByProperties',
                 'loadUnchanged',
-                'getQuery',
             ],
             true
         );
@@ -65,29 +58,6 @@ public function getTypeFromMethodCall(
             return ParametersAcceptorSelector::selectSingle($methodReflection->getVariants())->getReturnType();
         }
 
-        if ($methodReflection->getName() === 'getQuery') {
-            $returnType = ParametersAcceptorSelector::selectSingle($methodReflection->getVariants())->getReturnType();
-            if (!$returnType instanceof ObjectType) {
-                return $returnType;
-            }
-
-            if ((new ObjectType(ContentEntityStorageInterface::class))->isSuperTypeOf($callerType)->yes()) {
-                return new ContentEntityQueryType(
-                    $returnType->getClassName(),
-                    $returnType->getSubtractedType(),
-                    $returnType->getClassReflection()
-                );
-            }
-            if ((new ObjectType(ConfigEntityStorageInterface::class))->isSuperTypeOf($callerType)->yes()) {
-                return new ConfigEntityQueryType(
-                    $returnType->getClassName(),
-                    $returnType->getSubtractedType(),
-                    $returnType->getClassReflection()
-                );
-            }
-            return $returnType;
-        }
-
         if (!$callerType instanceof EntityStorageType) {
             $resolvedEntityType = $this->entityDataRepository->resolveFromStorage($callerType);
             if ($resolvedEntityType === null) {

src/Type/EntityStorage/GetQueryReturnTypeExtension.php

@@ -0,0 +1,61 @@
+<?php declare(strict_types=1);
+
+namespace mglaman\PHPStanDrupal\Type\EntityStorage;
+
+use Drupal\Core\Config\Entity\ConfigEntityStorageInterface;
+use Drupal\Core\Entity\ContentEntityStorageInterface;
+use Drupal\Core\Entity\EntityStorageInterface;
+use mglaman\PHPStanDrupal\Type\EntityQuery\ConfigEntityQueryType;
+use mglaman\PHPStanDrupal\Type\EntityQuery\ContentEntityQueryType;
+use PhpParser\Node\Expr\MethodCall;
+use PHPStan\Analyser\Scope;
+use PHPStan\Reflection\MethodReflection;
+use PHPStan\Reflection\ParametersAcceptorSelector;
+use PHPStan\Type\DynamicMethodReturnTypeExtension;
+use PHPStan\Type\ObjectType;
+
+final class GetQueryReturnTypeExtension implements DynamicMethodReturnTypeExtension
+{
+
+    public function getClass(): string
+    {
+        return EntityStorageInterface::class;
+    }
+
+    public function isMethodSupported(MethodReflection $methodReflection): bool
+    {
+        return $methodReflection->getName() === 'getQuery';
+    }
+
+    public function getTypeFromMethodCall(
+        MethodReflection $methodReflection,
+        MethodCall $methodCall,
+        Scope $scope
+    ): \PHPStan\Type\Type {
+        $returnType = ParametersAcceptorSelector::selectSingle($methodReflection->getVariants())->getReturnType();
+        if (!$returnType instanceof ObjectType) {
+            return $returnType;
+        }
+
+        $callerType = $scope->getType($methodCall->var);
+        if (!$callerType instanceof ObjectType) {
+            return $returnType;
+        }
+
+        if ((new ObjectType(ContentEntityStorageInterface::class))->isSuperTypeOf($callerType)->yes()) {
+            return new ContentEntityQueryType(
+                $returnType->getClassName(),
+                $returnType->getSubtractedType(),
+                $returnType->getClassReflection()
+            );
+        }
+        if ((new ObjectType(ConfigEntityStorageInterface::class))->isSuperTypeOf($callerType)->yes()) {
+            return new ConfigEntityQueryType(
+                $returnType->getClassName(),
+                $returnType->getSubtractedType(),
+                $returnType->getClassReflection()
+            );
+        }
+        return $returnType;
+    }
+}