Try a workaround to infer the type of the entity storage when it has been type hinted

Author: brambaud

src/Type/EntityStorage/EntityStorageDynamicReturnTypeExtension.php

@@ -16,6 +16,7 @@
 use PHPStan\Type\IntegerType;
 use PHPStan\Type\ObjectType;
 use PHPStan\Type\StringType;
+use PHPStan\Type\Type;
 use PHPStan\Type\TypeCombinator;
 
 class EntityStorageDynamicReturnTypeExtension implements DynamicMethodReturnTypeExtension
@@ -60,6 +61,18 @@ public function getTypeFromMethodCall(
         $callerType = $scope->getType($methodCall->var);
 
         if (!$callerType instanceof EntityStorageType) {
+            if (!$callerType instanceof ObjectType) {
+                return ParametersAcceptorSelector::selectSingle($methodReflection->getVariants())->getReturnType();
+            }
+
+            // Workaround because cannot figure out why the caller type is not an EntityStorageType
+            // when it has been type hinted.
+            // Instead, we try to infer the type, i.e. ContentEntityStorageType or ConfigEntityStorageType, here.
+            // @todo: we should definitively look for other cases that getQuery.
+            if ($methodReflection->getName() === 'getQuery') {
+                return $this->getReturnTypeForGetQueryMethod($methodReflection, EntityStorageHelper::getTypeFromStorageObjectType($callerType));
+            }
+
             return ParametersAcceptorSelector::selectSingle($methodReflection->getVariants())->getReturnType();
         }
 
@@ -79,27 +92,32 @@ 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 $this->getReturnTypeForGetQueryMethod($methodReflection, $callerType);
         }
 
         return $type;
     }
+
+    private function getReturnTypeForGetQueryMethod(MethodReflection $methodReflection, ?Type $entityStorageType): Type
+    {
+        $returnType = ParametersAcceptorSelector::selectSingle($methodReflection->getVariants())->getReturnType();
+        if (!$returnType instanceof ObjectType) {
+            return $returnType;
+        }
+        if ($entityStorageType instanceof ContentEntityStorageType) {
+            return new ContentEntityQueryType(
+                $returnType->getClassName(),
+                $returnType->getSubtractedType(),
+                $returnType->getClassReflection()
+            );
+        }
+        if ($entityStorageType instanceof ConfigEntityStorageType) {
+            return new ConfigEntityQueryType(
+                $returnType->getClassName(),
+                $returnType->getSubtractedType(),
+                $returnType->getClassReflection()
+            );
+        }
+        return $returnType;
+    }
 }

src/Type/EntityStorage/EntityStorageHelper.php

@@ -0,0 +1,39 @@
+<?php
+
+declare(strict_types=1);
+
+namespace mglaman\PHPStanDrupal\Type\EntityStorage;
+
+use Drupal\Core\Config\Entity\ConfigEntityStorage;
+use Drupal\Core\Config\Entity\ConfigEntityStorageInterface;
+use Drupal\Core\Entity\ContentEntityStorageInterface;
+use Drupal\Core\Entity\EntityStorageInterface;
+use Drupal\Core\Entity\Sql\SqlContentEntityStorage;
+use Drupal\Core\Entity\Sql\SqlEntityStorageInterface;
+use PHPStan\Type\ObjectType;
+
+final class EntityStorageHelper
+{
+    public static function getTypeFromStorageObjectType(ObjectType $classType): ?EntityStorageType
+    {
+        if ((new ObjectType(ConfigEntityStorageInterface::class))->isSuperTypeOf($classType)->yes()) {
+            $storageClassName = ConfigEntityStorage::class;
+        } elseif ((new ObjectType(SqlEntityStorageInterface::class))->isSuperTypeOf($classType)->yes()) {
+            $storageClassName = SqlContentEntityStorage::class;
+        } else {
+            return null;
+        }
+
+        $storageType = new ObjectType($storageClassName);
+        if ((new ObjectType(EntityStorageInterface::class))->isSuperTypeOf($storageType)->no()) {
+            return null;
+        }
+        if ((new ObjectType(ConfigEntityStorageInterface::class))->isSuperTypeOf($storageType)->yes()) {
+            return new ConfigEntityStorageType('', $storageClassName);
+        }
+        if ((new ObjectType(ContentEntityStorageInterface::class))->isSuperTypeOf($storageType)->yes()) {
+            return new ContentEntityStorageType('', $storageClassName);
+        }
+        return new EntityStorageType('', $storageClassName);
+    }
+}