Move getQuery method return type to own extension

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

@@ -4,17 +4,13 @@
 
 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\TypeCombinator;
 
@@ -46,7 +42,6 @@ public function isMethodSupported(MethodReflection $methodReflection): bool
                 'loadMultiple',
                 'loadByProperties',
                 'loadUnchanged',
-                'getQuery',
             ],
             true
         );
@@ -78,27 +73,6 @@ public function getTypeFromMethodCall(
 
             return new ArrayType(new IntegerType(), $type);
         }
-        if ($methodReflection->getName() === 'getQuery') {
-            $returnType = ParametersAcceptorSelector::selectSingle($methodReflection->getVariants())->getReturnType();
-            if (!$returnType instanceof ObjectType) {
-                return $returnType;
-            }
-            if ($callerType instanceof ContentEntityStorageType) {
-                return new ContentEntityQueryType(
-                    $returnType->getClassName(),
-                    $returnType->getSubtractedType(),
-                    $returnType->getClassReflection()
-                );
-            }
-            if ($callerType instanceof ConfigEntityStorageType) {
-                return new ConfigEntityQueryType(
-                    $returnType->getClassName(),
-                    $returnType->getSubtractedType(),
-                    $returnType->getClassReflection()
-                );
-            }
-            return $returnType;
-        }
 
         return $type;
     }

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;
+    }
+}