From 4ed6cf236e58144757a3f9a30a45797ff66d9f2d Mon Sep 17 00:00:00 2001
From: IanDelMar <42134098+IanDelMar@users.noreply.github.com>
Date: Sat, 14 Oct 2023 20:50:24 +0200
Subject: [PATCH] Add dynamic return type extension for `wp_die()` (#201)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Add dynamic return type extension for wp_die()

* Update tests for wp_die()

Co-authored-by: Viktor Szépe <viktor@szepe.net>

---------

Co-authored-by: Viktor Szépe <viktor@szepe.net>
---
 extension.neon                                |  5 +-
 ...pDieDynamicFunctionReturnTypeExtension.php | 53 +++++++++++++++++++
 tests/DynamicReturnTypeExtensionTest.php      |  1 +
 tests/data/wp_die.php                         | 15 ++++++
 4 files changed, 73 insertions(+), 1 deletion(-)
 create mode 100644 src/WpDieDynamicFunctionReturnTypeExtension.php
 create mode 100644 tests/data/wp_die.php

diff --git a/extension.neon b/extension.neon
index 5d2a70a..39e0472 100644
--- a/extension.neon
+++ b/extension.neon
@@ -95,6 +95,10 @@ services:
         class: SzepeViktor\PHPStan\WordPress\WpParseUrlFunctionDynamicReturnTypeExtension
         tags:
             - phpstan.broker.dynamicFunctionReturnTypeExtension
+    -
+        class: SzepeViktor\PHPStan\WordPress\WpDieDynamicFunctionReturnTypeExtension
+        tags:
+            - phpstan.broker.dynamicFunctionReturnTypeExtension
     -
         class: SzepeViktor\PHPStan\WordPress\HookDocsVisitor
         tags:
@@ -124,7 +128,6 @@ parameters:
         - SAVEQUERIES
         - SCRIPT_DEBUG
     earlyTerminatingFunctionCalls:
-        - wp_die
         - wp_send_json
         - wp_send_json_success
         - wp_send_json_error
diff --git a/src/WpDieDynamicFunctionReturnTypeExtension.php b/src/WpDieDynamicFunctionReturnTypeExtension.php
new file mode 100644
index 0000000..bce6590
--- /dev/null
+++ b/src/WpDieDynamicFunctionReturnTypeExtension.php
@@ -0,0 +1,53 @@
+<?php
+
+/**
+ * Set return type of wp_die().
+ */
+
+declare(strict_types=1);
+
+namespace SzepeViktor\PHPStan\WordPress;
+
+use PhpParser\Node\Expr\FuncCall;
+use PHPStan\Analyser\Scope;
+use PHPStan\Reflection\FunctionReflection;
+use PHPStan\Type\Constant\ConstantStringType;
+use PHPStan\Type\Type;
+use PHPStan\Type\VoidType;
+use PHPStan\Type\NeverType;
+
+class WpDieDynamicFunctionReturnTypeExtension implements \PHPStan\Type\DynamicFunctionReturnTypeExtension
+{
+    public function isFunctionSupported(FunctionReflection $functionReflection): bool
+    {
+        return $functionReflection->getName() === 'wp_die';
+    }
+
+    // phpcs:ignore SlevomatCodingStandard.Functions.UnusedParameter
+    public function getTypeFromFunctionCall(FunctionReflection $functionReflection, FuncCall $functionCall, Scope $scope): Type
+    {
+        $args = $functionCall->getArgs();
+
+        // Called without $args parameter
+        if (count($args) < 3) {
+            return new NeverType();
+        }
+
+        $argType = $scope->getType($args[2]->value);
+
+        // Return void for non constant arrays.
+        if (! $argType->isConstantArray()->yes()) {
+            return new VoidType();
+        }
+
+        // Return never if the key 'exit' is not set.
+        if (! $argType->hasOffsetValueType(new ConstantStringType('exit'))->yes()) {
+            return new NeverType();
+        }
+
+        // Note WP's wp_die handlers do lazy comparison
+        return $argType->getOffsetValueType(new ConstantStringType('exit'))->toBoolean()->isTrue()->yes()
+            ? new NeverType()
+            : new VoidType();
+    }
+}
diff --git a/tests/DynamicReturnTypeExtensionTest.php b/tests/DynamicReturnTypeExtensionTest.php
index 92ea1a1..69471ae 100644
--- a/tests/DynamicReturnTypeExtensionTest.php
+++ b/tests/DynamicReturnTypeExtensionTest.php
@@ -31,6 +31,7 @@ public function dataFileAsserts(): iterable
         yield from $this->gatherAssertTypes(__DIR__ . '/data/term_exists.php');
         yield from $this->gatherAssertTypes(__DIR__ . '/data/wp_error_parameter.php');
         yield from $this->gatherAssertTypes(__DIR__ . '/data/wp_parse_url.php');
+        yield from $this->gatherAssertTypes(__DIR__ . '/data/wp_die.php');
         yield from $this->gatherAssertTypes(__DIR__ . '/data/wp_theme_get.php');
     }
 
diff --git a/tests/data/wp_die.php b/tests/data/wp_die.php
new file mode 100644
index 0000000..0f2744f
--- /dev/null
+++ b/tests/data/wp_die.php
@@ -0,0 +1,15 @@
+<?php
+
+declare(strict_types=1);
+
+namespace SzepeViktor\PHPStan\WordPress\Tests;
+
+use function PHPStan\Testing\assertType;
+
+/** @var array $array */
+$array = null;
+
+assertType('*NEVER*', wp_die('', ''));
+assertType('*NEVER*', wp_die('', '', ['exit' => true]));
+assertType('void', wp_die('', '', ['exit' => false]));
+assertType('void', wp_die('', '', $array));