[code-quality] add with() + callback() support to TypeWillReturnCallableArrowFunctionRector (#535)

* add fixture to TypeWillReturnCallableArrowFunctionRector

* add with() + callback() support to TypeWillReturnCallableArrowFunctionRector

Author: TomasVotruba

rules-tests/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector/Fixture/with_callback.php.inc

@@ -0,0 +1,37 @@
+<?php
+
+namespace Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector\Fixture;
+
+use PHPUnit\Framework\TestCase;
+use Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector\Source\SomeMockedClass;
+
+final class FillKnownParamType extends TestCase
+{
+    public function test($value): void
+    {
+        $this->createMock(SomeMockedClass::class)
+            ->method('someMethod')
+            ->with($this->callback(fn ($name) => $value));
+    }
+}
+
+?>
+-----
+<?php
+
+namespace Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector\Fixture;
+
+use PHPUnit\Framework\TestCase;
+use Rector\PHPUnit\Tests\CodeQuality\Rector\Class_\TypeWillReturnCallableArrowFunctionRector\Source\SomeMockedClass;
+
+final class FillKnownParamType extends TestCase
+{
+    public function test($value): void
+    {
+        $this->createMock(SomeMockedClass::class)
+            ->method('someMethod')
+            ->with($this->callback(fn (string $name): int => $value));
+    }
+}
+
+?>

rules/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector.php

@@ -146,13 +146,8 @@ public function refactor(Node $node): ?Class_
                 return null;
             }
 
-            if (! $this->isName($node->name, self::WILL_RETURN_CALLBACK)) {
-                return null;
-            }
-
-            $innerArg = $node->getArgs()[0]
-                ->value;
-            if (! $innerArg instanceof ArrowFunction && ! $innerArg instanceof Closure) {
+            $innerClosure = $this->matchInnerClosure($node);
+            if (! $innerClosure instanceof Node) {
                 return null;
             }
 
@@ -205,7 +200,7 @@ public function refactor(Node $node): ?Class_
                 return null;
             }
 
-            foreach ($innerArg->params as $key => $param) {
+            foreach ($innerClosure->params as $key => $param) {
                 // avoid typing variadic parameters
                 if ($param->variadic) {
                     continue;
@@ -239,14 +234,14 @@ public function refactor(Node $node): ?Class_
                 $hasChanged = true;
             }
 
-            if (! $innerArg->returnType instanceof Node) {
+            if (! $innerClosure->returnType instanceof Node) {
                 $returnTypeNode = $this->staticTypeMapper->mapPHPStanTypeToPhpParserNode(
                     $parameterTypesAndReturnType->getReturnType(),
                     TypeKind::RETURN
                 );
 
                 if ($returnTypeNode instanceof Node) {
-                    $innerArg->returnType = $returnTypeNode;
+                    $innerClosure->returnType = $returnTypeNode;
                     $hasChanged = true;
                 }
             }
@@ -259,6 +254,33 @@ public function refactor(Node $node): ?Class_
         return $node;
     }
 
+    public function matchInnerClosure(MethodCall $methodCall): null|ArrowFunction|Closure
+    {
+        if ($this->isName($methodCall->name, 'with')) {
+            // special case for nested callback
+            $withFirstArg = $methodCall->getArgs()[0];
+
+            if ($withFirstArg->value instanceof MethodCall) {
+                $nestedMethodCall = $withFirstArg->value;
+                if ($this->isName($nestedMethodCall->name, 'callback')) {
+                    $nestedArg = $nestedMethodCall->getArgs()[0];
+                    if ($nestedArg->value instanceof ArrowFunction || $nestedArg->value instanceof Closure) {
+                        return $nestedArg->value;
+                    }
+                }
+            }
+        }
+
+        if ($this->isName($methodCall->name, self::WILL_RETURN_CALLBACK)) {
+            $innerArg = $methodCall->getArgs()[0];
+            if ($innerArg->value instanceof ArrowFunction || $innerArg->value instanceof Closure) {
+                return $innerArg->value;
+            }
+        }
+
+        return null;
+    }
+
     /**
      * @param array<string, string> $propertyNameToMockedTypes
      */