[TypeDeclaration] Allow Function_ and ArrowFunction on ReturnTypeFrom…

…ReturnNewRector (#583)
rectorphp · Aug 3, 2021 · ae55eac · ae55eac
1 parent 0888a92
commit ae55eac
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 4 deletions.
diff --git a/...laration/Rector/ClassMethod/ReturnTypeFromReturnNewRector/Fixture/arrow_function_.php.inc b/...laration/Rector/ClassMethod/ReturnTypeFromReturnNewRector/Fixture/arrow_function_.php.inc
@@ -0,0 +1,35 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Rector\Tests\TypeDeclaration\Rector\ClassMethod\ReturnTypeFromReturnNewRector\Fixture;
+
+use Rector\Tests\TypeDeclaration\Rector\ClassMethod\ReturnTypeFromReturnNewRector\Source\SomeResponse;
+
+final class SomeArrowFunction
+{
+    public function action()
+    {
+        return fn() => new SomeResponse();
+    }
+}
+
+?>
+-----
+<?php
+
+declare(strict_types=1);
+
+namespace Rector\Tests\TypeDeclaration\Rector\ClassMethod\ReturnTypeFromReturnNewRector\Fixture;
+
+use Rector\Tests\TypeDeclaration\Rector\ClassMethod\ReturnTypeFromReturnNewRector\Source\SomeResponse;
+
+final class SomeArrowFunction
+{
+    public function action()
+    {
+        return fn() : \Rector\Tests\TypeDeclaration\Rector\ClassMethod\ReturnTypeFromReturnNewRector\Source\SomeResponse => new SomeResponse();
+    }
+}
+
+?>
diff --git a/...ypeDeclaration/Rector/ClassMethod/ReturnTypeFromReturnNewRector/Fixture/function_.php.inc b/...ypeDeclaration/Rector/ClassMethod/ReturnTypeFromReturnNewRector/Fixture/function_.php.inc
@@ -0,0 +1,29 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Rector\Tests\TypeDeclaration\Rector\ClassMethod\ReturnTypeFromReturnNewRector\Fixture;
+
+use Rector\Tests\TypeDeclaration\Rector\ClassMethod\ReturnTypeFromReturnNewRector\Source\SomeResponse;
+
+function someFunction()
+{
+    return new SomeResponse();
+}
+
+?>
+-----
+<?php
+
+declare(strict_types=1);
+
+namespace Rector\Tests\TypeDeclaration\Rector\ClassMethod\ReturnTypeFromReturnNewRector\Fixture;
+
+use Rector\Tests\TypeDeclaration\Rector\ClassMethod\ReturnTypeFromReturnNewRector\Source\SomeResponse;
+
+function someFunction(): \Rector\Tests\TypeDeclaration\Rector\ClassMethod\ReturnTypeFromReturnNewRector\Source\SomeResponse
+{
+    return new SomeResponse();
+}
+
+?>
diff --git a/rules/TypeDeclaration/Rector/ClassMethod/ReturnTypeFromReturnNewRector.php b/rules/TypeDeclaration/Rector/ClassMethod/ReturnTypeFromReturnNewRector.php
@@ -5,9 +5,11 @@
 namespace Rector\TypeDeclaration\Rector\ClassMethod;
 
 use PhpParser\Node;
+use PhpParser\Node\Expr\ArrowFunction;
 use PhpParser\Node\Expr\New_;
 use PhpParser\Node\Name;
 use PhpParser\Node\Stmt\ClassMethod;
+use PhpParser\Node\Stmt\Function_;
 use PhpParser\Node\Stmt\Return_;
 use PHPStan\Type\ObjectType;
 use Rector\Core\Rector\AbstractRector;
@@ -60,20 +62,25 @@ public function action(): Response
      */
     public function getNodeTypes(): array
     {
-        return [ClassMethod::class];
+        return [ClassMethod::class, Function_::class, ArrowFunction::class];
     }
 
     /**
-     * @param ClassMethod $node
+     * @param ClassMethod|Function_|ArrowFunction $node
      */
     public function refactor(Node $node): ?Node
     {
         if ($node->returnType !== null) {
             return null;
         }
 
-        /** @var Return_[] $returns */
-        $returns = $this->betterNodeFinder->findInstanceOf((array) $node->stmts, Return_::class);
+        if ($node instanceof ArrowFunction) {
+            $returns = [new Return_($node->expr)];
+        } else {
+            /** @var Return_[] $returns */
+            $returns = $this->betterNodeFinder->findInstanceOf((array) $node->stmts, Return_::class);
+        }
+
         if ($returns === []) {
             return null;
         }