ObjectRepositoryDynamicReturnTypeExtension - do not override return types for existing native methods

Author: ondrejmirtes

Diff for: src/Type/Doctrine/ObjectRepositoryDynamicReturnTypeExtension.php

@@ -4,17 +4,29 @@
 
 use PhpParser\Node\Expr\MethodCall;
 use PHPStan\Analyser\Scope;
+use PHPStan\Broker\Broker;
+use PHPStan\Reflection\BrokerAwareExtension;
 use PHPStan\Reflection\MethodReflection;
+use PHPStan\Reflection\ParametersAcceptorSelector;
 use PHPStan\Type\ArrayType;
 use PHPStan\Type\IntegerType;
 use PHPStan\Type\MixedType;
 use PHPStan\Type\ObjectType;
 use PHPStan\Type\Type;
 use PHPStan\Type\TypeCombinator;
+use PHPStan\Type\TypeWithClassName;
 
-class ObjectRepositoryDynamicReturnTypeExtension implements \PHPStan\Type\DynamicMethodReturnTypeExtension
+class ObjectRepositoryDynamicReturnTypeExtension implements \PHPStan\Type\DynamicMethodReturnTypeExtension, BrokerAwareExtension
 {
 
+	/** @var Broker */
+	private $broker;
+
+	public function setBroker(Broker $broker): void
+	{
+		$this->broker = $broker;
+	}
+
 	public function getClass(): string
 	{
 		return 'Doctrine\Common\Persistence\ObjectRepository';
@@ -36,10 +48,35 @@ public function getTypeFromMethodCall(
 	): Type
 	{
 		$calledOnType = $scope->getType($methodCall->var);
-		if (!$calledOnType instanceof ObjectRepositoryType) {
+		if (!$calledOnType instanceof TypeWithClassName) {
 			return new MixedType();
 		}
+
 		$methodName = $methodReflection->getName();
+		$repositoryClassReflection = $this->broker->getClass($calledOnType->getClassName());
+		if (
+			(
+				(
+					strpos($methodName, 'findBy') === 0
+					&& strlen($methodName) > strlen('findBy')
+				) || (
+					strpos($methodName, 'findOneBy') === 0
+					&& strlen($methodName) > strlen('findOneBy')
+				)
+			)
+			&& $repositoryClassReflection->hasNativeMethod($methodName)
+		) {
+			return ParametersAcceptorSelector::selectFromArgs(
+				$scope,
+				$methodCall->args,
+				$repositoryClassReflection->getNativeMethod($methodName)->getVariants()
+			)->getReturnType();
+		}
+
+		if (!$calledOnType instanceof ObjectRepositoryType) {
+			return new MixedType();
+		}
+
 		$entityType = new ObjectType($calledOnType->getEntityClass());
 
 		if ($methodName === 'find' || strpos($methodName, 'findOneBy') === 0) {

Diff for: tests/DoctrineIntegration/ORM/data/customRepositoryUsage-2.json

@@ -13,5 +13,15 @@
         "message": "Call to an undefined method PHPStan\\DoctrineIntegration\\ORM\\CustomRepositoryUsage\\MyEntity::nonexistent().",
         "line": 47,
         "ignorable": true
+    },
+    {
+        "message": "Cannot call method test() on int.",
+        "line": 52,
+        "ignorable": true
+    },
+    {
+        "message": "Cannot call method test() on int.",
+        "line": 53,
+        "ignorable": true
     }
 ]

Diff for: tests/DoctrineIntegration/ORM/data/customRepositoryUsage.php

@@ -46,6 +46,12 @@ public function genericRepository(): void
 		$entity->doSomethingElse();
 		$entity->nonexistent();
 	}
+
+	public function callExistingMethodOnRepository(): void
+	{
+		$this->repository->findOneByBlabla()->test();
+		$this->anotherRepository->findOneByBlabla()->test();
+	}
 }
 
 /**