Add return type extension for echo key in args array

extension.neon

@@ -11,6 +11,10 @@ services:
         class: SzepeViktor\PHPStan\WordPress\WpThemeMagicPropertiesClassReflectionExtension
         tags:
             - phpstan.broker.propertiesClassReflectionExtension
+    -
+        class: SzepeViktor\PHPStan\WordPress\EchoKeyDynamicFunctionReturnTypeExtension
+        tags:
+            - phpstan.broker.dynamicFunctionReturnTypeExtension
     -
         class: SzepeViktor\PHPStan\WordPress\GetPermalinkDynamicFunctionReturnTypeExtension
         tags:

src/EchoKeyDynamicFunctionReturnTypeExtension.php

@@ -0,0 +1,160 @@
+<?php
+
+/**
+ * Set return type of various functions that support an optional `$args`
+ * of type array with `echo` key that defaults to true|1.
+ */
+
+declare(strict_types=1);
+
+namespace SzepeViktor\PHPStan\WordPress;
+
+use PhpParser\Node\Expr\FuncCall;
+use PHPStan\Analyser\Scope;
+use PHPStan\Reflection\FunctionReflection;
+use PHPStan\Reflection\ParametersAcceptorSelector;
+use PHPStan\Type\Constant\ConstantArrayType;
+use PHPStan\Type\Constant\ConstantBooleanType;
+use PHPStan\Type\Constant\ConstantIntegerType;
+use PHPStan\Type\Constant\ConstantStringType;
+use PHPStan\Type\NullType;
+use PHPStan\Type\Type;
+use PHPStan\Type\TypeCombinator;
+use PHPStan\Type\VoidType;
+
+class EchoKeyDynamicFunctionReturnTypeExtension implements \PHPStan\Type\DynamicFunctionReturnTypeExtension
+{
+    /**
+     * Function name and position of `$args` parameter.
+     */
+    private const SUPPORTED_FUNCTIONS = [
+        'get_search_form' => 0,
+        'the_title_attribute' => 0,
+        'wp_dropdown_categories' => 0,
+        'wp_dropdown_languages' => 0,
+        'wp_dropdown_pages' => 0,
+        'wp_dropdown_users' => 0,
+        'wp_get_archives' => 0,
+        'wp_link_pages' => 0,
+        'wp_list_authors' => 0,
+        'wp_list_bookmarks' => 0,
+        'wp_list_categories' => 0,
+        'wp_list_comments' => 0,
+        'wp_list_pages' => 0,
+        'wp_list_users' => 0,
+        'wp_login_form' => 0,
+        'wp_page_menu' => 0,
+    ];
+
+    /**
+     * Functions with strictly boolean `echo` key.
+     */
+    private const STRICTLY_BOOL = [
+        'get_search_form',
+        'the_title_attribute',
+        'wp_list_authors',
+        'wp_list_comments',
+        'wp_list_pages',
+        'wp_list_users',
+        'wp_login_form',
+        'wp_page_menu',
+    ];
+
+    public function isFunctionSupported(FunctionReflection $functionReflection): bool
+    {
+        return array_key_exists($functionReflection->getName(), self::SUPPORTED_FUNCTIONS);
+    }
+
+    public function getTypeFromFunctionCall(FunctionReflection $functionReflection, FuncCall $functionCall, Scope $scope): ?Type
+    {
+        $name = $functionReflection->getName();
+        $functionParameter = self::SUPPORTED_FUNCTIONS[$name] ?? null;
+        $args = $functionCall->getArgs();
+
+        if ($functionParameter === null) {
+            throw new \PHPStan\ShouldNotHappenException(
+                sprintf(
+                    'Could not detect return types for function %s()',
+                    $name
+                )
+            );
+        }
+
+        if (!isset($args[$functionParameter])) {
+            return self::getEchoTrueReturnType($name);
+        }
+
+        $argumentType = $scope->getType($args[$functionParameter]->value);
+        $echoType = self::getEchoType($argumentType);
+        $defaultType = ParametersAcceptorSelector::selectFromArgs(
+            $scope,
+            $functionCall->getArgs(),
+            $functionReflection->getVariants()
+        )->getReturnType();
+
+        if ($echoType instanceof ConstantBooleanType) {
+            return ($echoType->getValue() === false)
+                ? self::maybeRemoveVoid($name, $defaultType)
+                : self::getEchoTrueReturnType($name);
+        }
+
+        if (!in_array($name, self::STRICTLY_BOOL, true) && $echoType instanceof ConstantIntegerType) {
+            return ($echoType->getValue() === 0)
+                ? self::maybeRemoveVoid($name, $defaultType)
+                : self::getEchoTrueReturnType($name);
+        }
+
+        return TypeCombinator::union($defaultType, new VoidType());
+    }
+
+    protected static function getEchoType(Type $argumentType): Type
+    {
+        $echoType = new ConstantBooleanType(true);
+
+        if ($argumentType instanceof ConstantArrayType) {
+            foreach ($argumentType->getKeyTypes() as $index => $key) {
+                if (! $key instanceof ConstantStringType || $key->getValue() !== 'echo') {
+                    continue;
+                }
+                $echoType = $argumentType->getValueTypes()[$index];
+            }
+        }
+        return $echoType;
+    }
+
+    protected static function maybeRemoveVoid(string $name, Type $type): Type
+    {
+        // These function can return void even if echo is not true/truthy.
+        $doNotRemove = [
+            'the_title_attribute',
+            'wp_dropdown_languages',
+            'wp_get_archives',
+            'wp_list_comments',
+        ];
+
+        // Fix omitted void type in WP doc block.
+        $type = TypeCombinator::union($type, new VoidType());
+
+        if ($name === 'wp_list_users') {
+            // null instead of void in WP doc block.
+            $type = TypeCombinator::remove($type, new NullType());
+        }
+
+        if (!in_array($name, $doNotRemove, true)) {
+            $type = TypeCombinator::remove($type, new VoidType());
+        }
+        return $type;
+    }
+
+    protected static function getEchoTrueReturnType( string $name ): Type
+    {
+        $type = [
+            'wp_list_categories' => TypeCombinator::union(
+                new VoidType(),
+                new ConstantBooleanType(false)
+            ),
+        ];
+
+        return $type[$name] ?? new VoidType();
+    }
+}

tests/DynamicReturnTypeExtensionTest.php

@@ -15,6 +15,7 @@ public function dataFileAsserts(): iterable
         yield from $this->gatherAssertTypes(__DIR__ . '/data/_get_list_table.php');
         yield from $this->gatherAssertTypes(__DIR__ . '/data/apply_filters.php');
         yield from $this->gatherAssertTypes(__DIR__ . '/data/current_time.php');
+        yield from $this->gatherAssertTypes(__DIR__ . '/data/echo_key.php');
         yield from $this->gatherAssertTypes(__DIR__ . '/data/echo_parameter.php');
         yield from $this->gatherAssertTypes(__DIR__ . '/data/esc_sql.php');
         yield from $this->gatherAssertTypes(__DIR__ . '/data/get_comment.php');

tests/data/echo_key.php

@@ -0,0 +1,74 @@
+<?php
+
+declare(strict_types=1);
+
+namespace SzepeViktor\PHPStan\WordPress\Tests;
+
+use function PHPStan\Testing\assertType;
+
+// Default value of true
+assertType('void', get_search_form());
+assertType('void', the_title_attribute());
+assertType('void', wp_dropdown_categories());
+assertType('void', wp_dropdown_languages());
+assertType('void', wp_dropdown_pages());
+assertType('void', wp_dropdown_users());
+assertType('void', wp_get_archives());
+assertType('void', wp_link_pages());
+assertType('void', wp_list_authors());
+assertType('void', wp_list_bookmarks());
+assertType('void|false', wp_list_categories());
+assertType('void', wp_list_comments());
+assertType('void', wp_list_pages());
+assertType('void', wp_list_users());
+assertType('void', wp_login_form());
+assertType('void', wp_page_menu());
+
+// Explicit value of true
+$args = ['echo' => true];
+assertType('void', get_search_form($args));
+assertType('void', the_title_attribute($args));
+assertType('void', wp_dropdown_categories($args));
+assertType('void', wp_dropdown_languages($args));
+assertType('void', wp_dropdown_pages($args));
+assertType('void', wp_dropdown_users($args));
+assertType('void', wp_get_archives($args));
+assertType('void', wp_link_pages($args));
+assertType('void', wp_list_authors($args));
+assertType('void', wp_list_bookmarks($args));
+assertType('void|false', wp_list_categories($args));
+assertType('void', wp_list_comments($args));
+assertType('void', wp_list_pages($args));
+assertType('void', wp_list_users($args));
+assertType('void', wp_login_form($args));
+assertType('void', wp_page_menu($args));
+
+// Explicit value of 1
+$args = ['echo' => 1];
+assertType('void', wp_dropdown_categories($args));
+assertType('void', wp_dropdown_languages($args));
+assertType('void', wp_dropdown_pages($args));
+assertType('void', wp_dropdown_users($args));
+assertType('void', wp_get_archives($args));
+assertType('void', wp_link_pages($args));
+assertType('void', wp_list_bookmarks($args));
+assertType('void|false', wp_list_categories($args));
+
+// Explicit value of false
+$args = ['echo' => false];
+assertType('string', get_search_form($args));
+assertType('string|void', the_title_attribute($args));
+assertType('string', wp_dropdown_categories($args));
+assertType('string|void', wp_dropdown_languages($args));
+assertType('string', wp_dropdown_pages($args));
+assertType('string', wp_dropdown_users($args));
+assertType('string|void', wp_get_archives($args));
+assertType('string', wp_link_pages($args));
+assertType('string', wp_list_authors($args));
+assertType('string', wp_list_bookmarks($args));
+assertType('string|false', wp_list_categories($args));
+assertType('string|void', wp_list_comments($args));
+assertType('string', wp_list_pages($args));
+assertType('string', wp_list_users($args));
+assertType('string', wp_login_form($args));
+assertType('string', wp_page_menu($args));