add with() + callback() support to TypeWillReturnCallableArrowFunctionRector

Author: TomasVotruba

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
      */