add with() + callback() support to TypeWillReturnCallableArrowFunctionRector

Author: TomasVotruba

rules/CodeQuality/Rector/Class_/TypeWillReturnCallableArrowFunctionRector.php

@@ -146,13 +146,32 @@ public function refactor(Node $node): ?Class_
                 return null;
             }
 
-            if (! $this->isName($node->name, self::WILL_RETURN_CALLBACK)) {
+            $innerClosure = null;
+
+            if ($this->isName($node->name, 'with')) {
+                // special case for nested callback
+                $withFirstArg = $node->getArgs()[0];
+
+                if ($withFirstArg->value instanceof MethodCall) {
+                    $nestedMethodCall = $withFirstArg->value;
+                    if ($this->isName($nestedMethodCall->name, 'callback')) {
+                        $innerClosure = $nestedMethodCall->getArgs()[0]
+                            ->value;
+                    }
+                }
+            } elseif ($this->isName($node->name, self::WILL_RETURN_CALLBACK)) {
+                $innerArg = $node->getArgs()[0];
+                if (! $innerArg->value instanceof ArrowFunction && ! $innerArg->value instanceof Closure) {
+                    return null;
+                }
+
+                $innerClosure = $innerArg->value;
+            } else {
                 return null;
             }
 
-            $innerArg = $node->getArgs()[0]
-                ->value;
-            if (! $innerArg instanceof ArrowFunction && ! $innerArg instanceof Closure) {
+            // no match
+            if (! $innerClosure instanceof Node\FunctionLike) {
                 return null;
             }
 
@@ -205,7 +224,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 +258,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;
                 }
             }