From e644df734e1bbe95810e0f617d17df091048a94e Mon Sep 17 00:00:00 2001
From: Pascal Birchler <pascalb@google.com>
Date: Thu, 22 Sep 2022 15:14:50 +0200
Subject: [PATCH] Add `get_sites()` dynamic function return type extension
 (#117)

* Add tests for `get_posts()`

* Add `get_sites()` dynamic function return type extension

* Simplify switch

* Simplify `isFunctionSupported`

* No null coalescing
---
 extension.neon                                |  4 +
 ...ostsDynamicFunctionReturnTypeExtension.php |  2 +-
 ...itesDynamicFunctionReturnTypeExtension.php | 75 +++++++++++++++++++
 tests/DynamicReturnTypeExtensionTest.php      |  2 +
 tests/data/get_posts.php                      | 12 +++
 tests/data/get_sites.php                      | 11 +++
 6 files changed, 105 insertions(+), 1 deletion(-)
 create mode 100644 src/GetSitesDynamicFunctionReturnTypeExtension.php
 create mode 100644 tests/data/get_posts.php
 create mode 100644 tests/data/get_sites.php

diff --git a/extension.neon b/extension.neon
index 427e107..c217aea 100644
--- a/extension.neon
+++ b/extension.neon
@@ -35,6 +35,10 @@ services:
         class: SzepeViktor\PHPStan\WordPress\GetPostsDynamicFunctionReturnTypeExtension
         tags:
             - phpstan.broker.dynamicFunctionReturnTypeExtension
+    -
+        class: SzepeViktor\PHPStan\WordPress\GetSitesDynamicFunctionReturnTypeExtension
+        tags:
+            - phpstan.broker.dynamicFunctionReturnTypeExtension
     -
         class: SzepeViktor\PHPStan\WordPress\GetTaxonomiesDynamicFunctionReturnTypeExtension
         tags:
diff --git a/src/GetPostsDynamicFunctionReturnTypeExtension.php b/src/GetPostsDynamicFunctionReturnTypeExtension.php
index 26ffb59..0a5b553 100644
--- a/src/GetPostsDynamicFunctionReturnTypeExtension.php
+++ b/src/GetPostsDynamicFunctionReturnTypeExtension.php
@@ -22,7 +22,7 @@ class GetPostsDynamicFunctionReturnTypeExtension implements \PHPStan\Type\Dynami
 {
     public function isFunctionSupported(FunctionReflection $functionReflection): bool
     {
-        return in_array($functionReflection->getName(), ['get_posts'], true);
+        return $functionReflection->getName() === 'get_posts';
     }
 
     /**
diff --git a/src/GetSitesDynamicFunctionReturnTypeExtension.php b/src/GetSitesDynamicFunctionReturnTypeExtension.php
new file mode 100644
index 0000000..aa5dbfd
--- /dev/null
+++ b/src/GetSitesDynamicFunctionReturnTypeExtension.php
@@ -0,0 +1,75 @@
+<?php
+
+/**
+ * Set return type of get_sites().
+ */
+
+declare(strict_types=1);
+
+namespace SzepeViktor\PHPStan\WordPress;
+
+use PhpParser\Node\Expr\FuncCall;
+use PHPStan\Analyser\Scope;
+use PHPStan\Reflection\FunctionReflection;
+use PHPStan\Type\Type;
+use PHPStan\Type\ArrayType;
+use PHPStan\Type\IntegerType;
+use PHPStan\Type\ObjectType;
+use PHPStan\Type\Constant\ConstantArrayType;
+use PHPStan\Type\Constant\ConstantStringType;
+
+class GetSitesDynamicFunctionReturnTypeExtension implements \PHPStan\Type\DynamicFunctionReturnTypeExtension
+{
+    public function isFunctionSupported(FunctionReflection $functionReflection): bool
+    {
+        return $functionReflection->getName() === 'get_sites';
+    }
+
+    /**
+     * @see https://developer.wordpress.org/reference/classes/wp_query/#return-fields-parameter
+     */
+    // phpcs:ignore SlevomatCodingStandard.Functions.UnusedParameter
+    public function getTypeFromFunctionCall(FunctionReflection $functionReflection, FuncCall $functionCall, Scope $scope): Type
+    {
+        $args = $functionCall->getArgs();
+
+        // Called without arguments
+        if (count($args) === 0) {
+            return new ArrayType(new IntegerType(), new ObjectType('WP_Site'));
+        }
+
+        $fields = '';
+
+        $argumentType = $scope->getType($args[0]->value);
+
+        // Called with an array argument
+        if ($argumentType instanceof ConstantArrayType) {
+            foreach ($argumentType->getKeyTypes() as $index => $key) {
+                if (! $key instanceof ConstantStringType || $key->getValue() !== 'fields') {
+                    continue;
+                }
+
+                $fieldsType = $argumentType->getValueTypes()[$index];
+                if ($fieldsType instanceof ConstantStringType) {
+                    $fields = $fieldsType->getValue();
+                }
+                break;
+            }
+        }
+
+        // Called with a string argument
+        if ($argumentType instanceof ConstantStringType) {
+            parse_str($argumentType->getValue(), $variables);
+            $fields = $variables['fields'] ?? 'all';
+        }
+
+        switch ($fields) {
+            case 'count':
+                return new IntegerType();
+            case 'ids':
+                return new ArrayType(new IntegerType(), new IntegerType());
+            default:
+                return new ArrayType(new IntegerType(), new ObjectType('WP_Site'));
+        }
+    }
+}
diff --git a/tests/DynamicReturnTypeExtensionTest.php b/tests/DynamicReturnTypeExtensionTest.php
index 155c9c9..c8757e9 100644
--- a/tests/DynamicReturnTypeExtensionTest.php
+++ b/tests/DynamicReturnTypeExtensionTest.php
@@ -20,6 +20,8 @@ public function dataFileAsserts(): iterable
         yield from $this->gatherAssertTypes(__DIR__ . '/data/get_comment.php');
         yield from $this->gatherAssertTypes(__DIR__ . '/data/get_object_taxonomies.php');
         yield from $this->gatherAssertTypes(__DIR__ . '/data/get_post.php');
+        yield from $this->gatherAssertTypes(__DIR__ . '/data/get_posts.php');
+        yield from $this->gatherAssertTypes(__DIR__ . '/data/get_sites.php');
         yield from $this->gatherAssertTypes(__DIR__ . '/data/get_terms.php');
         yield from $this->gatherAssertTypes(__DIR__ . '/data/has_filter.php');
         yield from $this->gatherAssertTypes(__DIR__ . '/data/mysql2date.php');
diff --git a/tests/data/get_posts.php b/tests/data/get_posts.php
new file mode 100644
index 0000000..9de82be
--- /dev/null
+++ b/tests/data/get_posts.php
@@ -0,0 +1,12 @@
+<?php
+
+declare(strict_types=1);
+
+namespace SzepeViktor\PHPStan\WordPress\Tests;
+
+use function PHPStan\Testing\assertType;
+
+assertType('array<int, WP_Post>', get_posts());
+assertType('array<int, WP_Post>', get_posts(['fields' => 'all']));
+assertType('array<int, int>', get_posts(['fields' => 'ids']));
+assertType('array<int, int>', get_posts(['fields' => 'id=>parent']));
diff --git a/tests/data/get_sites.php b/tests/data/get_sites.php
new file mode 100644
index 0000000..fd03fd1
--- /dev/null
+++ b/tests/data/get_sites.php
@@ -0,0 +1,11 @@
+<?php
+
+declare(strict_types=1);
+
+namespace SzepeViktor\PHPStan\WordPress\Tests;
+
+use function PHPStan\Testing\assertType;
+
+assertType('array<int, WP_Site>', get_sites());
+assertType('int', get_sites(['fields' => 'count']));
+assertType('array<int, int>', get_sites(['fields' => 'ids']));