From ae55eac4c937caf15d9bb320672f038915bc6f80 Mon Sep 17 00:00:00 2001
From: Abdul Malik Ikhsan <samsonasik@gmail.com>
Date: Tue, 3 Aug 2021 18:53:45 +0700
Subject: [PATCH] [TypeDeclaration] Allow Function_ and ArrowFunction on
 ReturnTypeFromReturnNewRector (#583)

---
 .../Fixture/arrow_function_.php.inc           | 35 +++++++++++++++++++
 .../Fixture/function_.php.inc                 | 29 +++++++++++++++
 .../ReturnTypeFromReturnNewRector.php         | 15 +++++---
 3 files changed, 75 insertions(+), 4 deletions(-)
 create mode 100644 rules-tests/TypeDeclaration/Rector/ClassMethod/ReturnTypeFromReturnNewRector/Fixture/arrow_function_.php.inc
 create mode 100644 rules-tests/TypeDeclaration/Rector/ClassMethod/ReturnTypeFromReturnNewRector/Fixture/function_.php.inc

diff --git a/rules-tests/TypeDeclaration/Rector/ClassMethod/ReturnTypeFromReturnNewRector/Fixture/arrow_function_.php.inc b/rules-tests/TypeDeclaration/Rector/ClassMethod/ReturnTypeFromReturnNewRector/Fixture/arrow_function_.php.inc
new file mode 100644
index 00000000000..ddc66c926d6
--- /dev/null
+++ b/rules-tests/TypeDeclaration/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/rules-tests/TypeDeclaration/Rector/ClassMethod/ReturnTypeFromReturnNewRector/Fixture/function_.php.inc b/rules-tests/TypeDeclaration/Rector/ClassMethod/ReturnTypeFromReturnNewRector/Fixture/function_.php.inc
new file mode 100644
index 00000000000..700a2f669e0
--- /dev/null
+++ b/rules-tests/TypeDeclaration/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
index e1329322326..a6787ccbd48 100644
--- 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,11 +62,11 @@ 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
     {
@@ -72,8 +74,13 @@ public function refactor(Node $node): ?Node
             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;
         }