added RepositoryReturnTypeExtension to support propagating entity type

extension.neon

@@ -3,6 +3,10 @@ services:
 		class: Nextras\OrmPhpStan\Types\CollectionReturnTypeExtension
 		tags:
 			- phpstan.broker.dynamicMethodReturnTypeExtension
+	-
+		class: Nextras\OrmPhpStan\Types\RepositoryReturnTypeExtension
+		tags:
+			- phpstan.broker.dynamicMethodReturnTypeExtension
 	-
 		class: Nextras\OrmPhpStan\Rules\SetValueMethodRule
 		tags:

src/Types/RepositoryReturnTypeExtension.php

@@ -0,0 +1,153 @@
+<?php declare(strict_types = 1);
+
+namespace Nextras\OrmPhpStan\Types;
+
+use Nextras\Orm\Collection\ICollection;
+use Nextras\Orm\Entity\IEntity;
+use Nextras\Orm\Repository\IRepository;
+use Nextras\Orm\Repository\Repository;
+use PhpParser\Node;
+use PhpParser\Node\Expr\MethodCall;
+use PhpParser\NodeFinder;
+use PhpParser\NodeTraverser;
+use PhpParser\NodeVisitor\NameResolver;
+use PHPStan\Analyser\Scope;
+use PHPStan\Parser\Parser;
+use PHPStan\Reflection\MethodReflection;
+use PHPStan\Reflection\ParametersAcceptorSelector;
+use PHPStan\ShouldNotHappenException;
+use PHPStan\Type\Constant\ConstantArrayType;
+use PHPStan\Type\Constant\ConstantStringType;
+use PHPStan\Type\DynamicMethodReturnTypeExtension;
+use PHPStan\Type\IntegerType;
+use PHPStan\Type\IntersectionType;
+use PHPStan\Type\IterableType;
+use PHPStan\Type\ObjectType;
+use PHPStan\Type\Type;
+use PHPStan\Type\TypeCombinator;
+use PHPStan\Type\TypeWithClassName;
+
+
+class RepositoryReturnTypeExtension implements DynamicMethodReturnTypeExtension
+{
+	/** @var Parser */
+	private $parser;
+
+
+	public function __construct(Parser $parser)
+	{
+		$this->parser = $parser;
+	}
+
+
+	public function getClass(): string
+	{
+		return IRepository::class;
+	}
+
+
+	public function isMethodSupported(MethodReflection $methodReflection): bool
+	{
+		static $methods = [
+			'getBy',
+			'getById',
+			'findAll',
+			'findBy',
+			'findById',
+		];
+		return in_array($methodReflection->getName(), $methods, true);
+	}
+
+
+	public function getTypeFromMethodCall(
+		MethodReflection $methodReflection,
+		MethodCall $methodCall,
+		Scope $scope
+	): Type
+	{
+		$repository = $scope->getType($methodCall->var);
+		\assert($repository instanceof TypeWithClassName);
+
+		if ($repository->getClassName() === Repository::class) {
+			return ParametersAcceptorSelector::selectSingle($methodReflection->getVariants())->getReturnType();
+		}
+
+		$repositoryReflection = new \ReflectionClass($repository->getClassName());
+		$entityClassNameTypes = $this->parseEntityClassNameTypes(
+			$repositoryReflection->getFileName(),
+			$repository->getClassName(),
+			$scope
+		);
+
+		if ($entityClassNameTypes === null) {
+			$entityType = new ObjectType(IEntity::class);
+		} else {
+			assert($entityClassNameTypes instanceof ConstantArrayType);
+			$classNameType = $entityClassNameTypes->getFirstValueType();
+			assert($classNameType instanceof ConstantStringType);
+			$entityType = new ObjectType($classNameType->getValue());
+		}
+
+		static $collectionReturnMethods = [
+			'findAll',
+			'findBy',
+			'findById',
+		];
+
+		static $entityReturnMethods = [
+			'getBy',
+			'getById',
+		];
+
+		$methodName = $methodReflection->getName();
+		if (in_array($methodName, $collectionReturnMethods, true)) {
+			return new IntersectionType([
+				new ObjectType(ICollection::class),
+				new IterableType(new IntegerType(), $entityType),
+			]);
+		} elseif (in_array($methodName, $entityReturnMethods, true)) {
+			return TypeCombinator::addNull($entityType);
+		}
+
+		throw new ShouldNotHappenException();
+	}
+
+
+	private function parseEntityClassNameTypes(string $fileName, string $className, Scope $scope): ?Type
+	{
+		$ast = $this->parser->parseFile($fileName);
+
+		$nodeTraverser = new NodeTraverser();
+		$nodeTraverser->addVisitor(new NameResolver());
+		$ast = $nodeTraverser->traverse($ast);
+
+		$nodeFinder = new NodeFinder();
+		$class = $nodeFinder->findFirst($ast, function (Node $node) use ($className) {
+			return $node instanceof Node\Stmt\Class_
+				&& $node->namespacedName->toString() === $className;
+		});
+
+		if ($class === null) {
+			return null;
+		}
+
+		$method = $nodeFinder->findFirst($class, function (Node $node) {
+			return $node instanceof Node\Stmt\ClassMethod
+				&& $node->name->name === 'getEntityClassNames';
+		});
+
+		if ($method === null) {
+			return null;
+		}
+
+		$return = $nodeFinder->findFirst($method, function (Node $node) {
+			return $node instanceof Node\Stmt\Return_;
+		});
+
+		if ($return instanceof Node\Stmt\Return_) {
+			return $scope->getType($return->expr);
+		} else {
+			return null;
+		}
+	}
+}

tests/expected.txt

@@ -6,3 +6,5 @@
 /testbox/Rules/Test.php:26:Entity NextrasTests\OrmPhpStan\Rules\Entity: property $age (int) does not accept string.
 /testbox/Types/CollectionTypesTest.php:15:Parameter #1 $author of method NextrasTests\OrmPhpStan\Types\CollectionTypesTest::takeAuthor() expects NextrasTests\OrmPhpStan\Types\Author, NextrasTests\OrmPhpStan\Types\Author|null given.
 /testbox/Types/CollectionTypesTest.php:16:Parameter #1 $author of method NextrasTests\OrmPhpStan\Types\CollectionTypesTest::takeAuthor() expects NextrasTests\OrmPhpStan\Types\Author, NextrasTests\OrmPhpStan\Types\Author|null given.
+/testbox/Types/RepositoryTypesTest.php:9:Parameter #1 $author of method NextrasTests\OrmPhpStan\Types\RepositoryTypesTest::takeAuthor() expects NextrasTests\OrmPhpStan\Types\Author, NextrasTests\OrmPhpStan\Types\Author|null given.
+/testbox/Types/RepositoryTypesTest.php:10:Parameter #1 $author of method NextrasTests\OrmPhpStan\Types\RepositoryTypesTest::takeAuthor() expects NextrasTests\OrmPhpStan\Types\Author, NextrasTests\OrmPhpStan\Types\Author|null given.

tests/testbox/Types/RepositoryTypesTest.php

@@ -0,0 +1,41 @@
+<?php declare(strict_types = 1);
+
+namespace NextrasTests\OrmPhpStan\Types;
+
+class RepositoryTypesTest
+{
+	public function testError(AuthorsRepository $repository)
+	{
+		$this->takeAuthor($repository->getById(1));
+		$this->takeAuthor($repository->getBy([]));
+	}
+
+
+	public function testOk(AuthorsRepository $repository)
+	{
+		$this->takeAuthorNullable($repository->getById(1));
+		$this->takeAuthorNullable($repository->getBy([]));
+		$this->takeAuthorNullable($repository->findAll()->fetch());
+		$this->takeAuthorNullable($repository->findBy([])->getById(1));
+		$this->takeAuthorNullable($repository->findBy([])->orderBy([])->limitBy(2, 0)->getById(1));
+		$this->takeAuthorArray($repository->findAll()->fetchAll());
+	}
+
+
+	private function takeAuthor(Author $author)
+	{
+	}
+
+
+	private function takeAuthorNullable(?Author $author)
+	{
+	}
+
+
+	/**
+	 * @param array<int, Author> $authors
+	 */
+	private function takeAuthorArray($authors)
+	{
+	}
+}

tests/testbox/Types/fixtures/AuthorsRepository.php

@@ -0,0 +1,14 @@
+<?php declare(strict_types = 1);
+
+namespace NextrasTests\OrmPhpStan\Types;
+
+use Nextras\Orm\Repository\Repository;
+
+
+class AuthorsRepository extends Repository
+{
+	public static function getEntityClassNames(): array
+	{
+		return [Author::class];
+	}
+}