Add EntityRepositoryInterface return type extension

extension.neon

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

src/Type/EntityRepositoryReturnTypeExtension.php

@@ -0,0 +1,96 @@
+<?php
+
+declare(strict_types=1);
+
+namespace mglaman\PHPStanDrupal\Type;
+
+use Drupal\Core\Config\Entity\ConfigEntityInterface;
+use Drupal\Core\Entity\EntityRepositoryInterface;
+use mglaman\PHPStanDrupal\Drupal\EntityDataRepository;
+use PhpParser\Node\Expr\MethodCall;
+use PHPStan\Analyser\Scope;
+use PHPStan\Reflection\MethodReflection;
+use PHPStan\Reflection\ParametersAcceptorSelector;
+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;
+
+final class EntityRepositoryReturnTypeExtension implements DynamicMethodReturnTypeExtension
+{
+
+    /**
+     * @var EntityDataRepository
+     */
+    private $entityDataRepository;
+
+    public function __construct(EntityDataRepository $entityDataRepository)
+    {
+        $this->entityDataRepository = $entityDataRepository;
+    }
+
+    public function getClass(): string
+    {
+        return EntityRepositoryInterface::class;
+    }
+
+    public function isMethodSupported(MethodReflection $methodReflection): bool
+    {
+        return in_array(
+            $methodReflection->getName(),
+            [
+                'getTranslationFromContext',
+                'loadEntityByUuid',
+                'loadEntityByConfigTarget',
+                'getActive',
+                'getActiveMultiple',
+                'getCanonical',
+                'getCanonicalMultiple',
+            ],
+            true
+        );
+    }
+
+    public function getTypeFromMethodCall(
+        MethodReflection $methodReflection,
+        MethodCall $methodCall,
+        Scope $scope
+    ): ?Type {
+        $methodName = $methodReflection->getName();
+        $methodArgs = $methodCall->getArgs();
+        $returnType = ParametersAcceptorSelector::selectSingle($methodReflection->getVariants())->getReturnType();
+
+        if (count($methodArgs) === 0) {
+            return $returnType;
+        }
+
+        if ($methodName === 'getTranslationFromContext') {
+            return $scope->getType($methodArgs[0]->value);
+        }
+
+        $entityObjectTypes = [];
+        $entityIdArg = $scope->getType($methodArgs[0]->value);
+        foreach ($entityIdArg->getConstantStrings() as $constantStringType) {
+            $entityObjectTypes[] = $this->entityDataRepository->get($constantStringType->getValue())->getClassType() ?? $returnType;
+        }
+        $entityTypes = TypeCombinator::union(...$entityObjectTypes);
+
+        if ($returnType->isArray()->no()) {
+            if ($returnType->isNull()->maybe()) {
+                $entityTypes = TypeCombinator::addNull($entityTypes);
+            }
+            return $entityTypes;
+        }
+
+        if ((new ObjectType(ConfigEntityInterface::class))->isSuperTypeOf($entityTypes)->yes()) {
+            $keyType = new StringType();
+        } else {
+            $keyType = new IntegerType();
+        }
+
+        return new ArrayType($keyType, $entityTypes);
+    }
+}

tests/src/Type/EntityRepositoryTypeTest.php

@@ -0,0 +1,33 @@
+<?php
+
+declare(strict_types=1);
+
+namespace mglaman\PHPStanDrupal\Tests\Type;
+
+use mglaman\PHPStanDrupal\Tests\AdditionalConfigFilesTrait;
+use PHPStan\Testing\TypeInferenceTestCase;
+
+final class EntityRepositoryTypeTest extends TypeInferenceTestCase
+{
+    use AdditionalConfigFilesTrait;
+
+    public function dataFileAsserts(): iterable
+    {
+        yield from self::gatherAssertTypes(__DIR__ . '/data/entity-repository.php');
+    }
+
+    /**
+     * @dataProvider dataFileAsserts
+     * @param string $assertType
+     * @param string $file
+     * @param mixed ...$args
+     */
+    public function testFileAsserts(
+        string $assertType,
+        string $file,
+        ...$args
+    ): void
+    {
+        $this->assertFileAsserts($assertType, $file, ...$args);
+    }
+}

tests/src/Type/data/entity-repository.php

@@ -0,0 +1,47 @@
+<?php
+
+namespace EntityRepository;
+
+use Drupal\node\Entity\Node;
+use function PHPStan\Testing\assertType;
+
+$entityRepository = \Drupal::service('entity.repository');
+
+assertType(
+    'Drupal\node\Entity\Node|null',
+    $entityRepository->loadEntityByUuid('node', '3f205175-04f7-4f57-b48b-9799299252c3')
+);
+
+assertType(
+    'Drupal\Core\Entity\Entity\EntityViewMode|null',
+    $entityRepository->loadEntityByConfigTarget('entity_view_mode', 'media.default')
+);
+
+assertType(
+    'Drupal\node\Entity\Node',
+    $entityRepository->getTranslationFromContext(Node::create())
+);
+
+assertType(
+    'Drupal\node\Entity\Node|null',
+    $entityRepository->getActive('node', 5)
+);
+
+assertType(
+    'array<int, Drupal\node\Entity\Node>',
+    $entityRepository->getActiveMultiple('node', [5])
+);
+assertType(
+    'array<string, Drupal\block\Entity\Block>',
+    $entityRepository->getActiveMultiple('block', ['foo'])
+);
+
+assertType(
+    'Drupal\node\Entity\Node|null',
+    $entityRepository->getCanonical('node', 5)
+);
+
+assertType(
+    'array<int, Drupal\node\Entity\Node>',
+    $entityRepository->getCanonicalMultiple('node', [5])
+);