Adjust logic when returning entity count type with access

Author: mglaman

src/Type/EntityQuery/EntityQueryCountType.php

@@ -2,8 +2,6 @@
 
 namespace mglaman\PHPStanDrupal\Type\EntityQuery;
 
-use PHPStan\Type\ObjectType;
-
 /**
  * Type used to represent an entity query instance as count query.
  */

src/Type/EntityQuery/EntityQueryDynamicReturnTypeExtension.php

@@ -39,29 +39,21 @@ public function getTypeFromMethodCall(
         $varType = $scope->getType($methodCall->var);
         $methodName = $methodReflection->getName();
 
-        if ($methodName === 'count') {
-            if ($varType instanceof EntityQueryType) {
-                $returnType = new EntityQueryCountType(
-                    $varType->getClassName(),
-                    $varType->getSubtractedType(),
-                    $varType->getClassReflection()
-                );
-
-                if ($varType->hasAccessCheck()) {
-                    return $returnType->withAccessCheck();
-                }
+        if (!$varType instanceof ObjectType) {
+            return $defaultReturnType;
+        }
 
-                return $returnType;
+        if ($methodName === 'count') {
+            $returnType = new EntityQueryCountType(
+                $varType->getClassName(),
+                $varType->getSubtractedType(),
+                $varType->getClassReflection()
+            );
+            if ($varType instanceof EntityQueryType && $varType->hasAccessCheck()) {
+                return $returnType->withAccessCheck();
             }
 
-            if ($varType instanceof ObjectType) {
-                return new EntityQueryCountType(
-                    $varType->getClassName(),
-                    $varType->getSubtractedType(),
-                    $varType->getClassReflection()
-                );
-            }
-            return $defaultReturnType;
+            return $returnType;
         }
 
         if ($methodName === 'execute') {