From a9b20179b169c97c3038067a3b7d2ef0e1c2b5f7 Mon Sep 17 00:00:00 2001
From: Toon Verwerft <toonverwerft@gmail.com>
Date: Fri, 9 Apr 2021 21:07:34 +0200
Subject: [PATCH] [Regex] add capture_groups() return type provider (#5)

---
 .../FunctionReturnTypeProvider.php            | 78 +++++++++++++++++++
 src/Plugin.php                                |  3 +
 2 files changed, 81 insertions(+)
 create mode 100644 src/EventHandler/Regex/CaptureGroups/FunctionReturnTypeProvider.php

diff --git a/src/EventHandler/Regex/CaptureGroups/FunctionReturnTypeProvider.php b/src/EventHandler/Regex/CaptureGroups/FunctionReturnTypeProvider.php
new file mode 100644
index 0000000..ff0e953
--- /dev/null
+++ b/src/EventHandler/Regex/CaptureGroups/FunctionReturnTypeProvider.php
@@ -0,0 +1,78 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Psl\Psalm\EventHandler\Regex\CaptureGroups;
+
+use Psalm\Plugin\EventHandler\Event\FunctionReturnTypeProviderEvent;
+use Psalm\Plugin\EventHandler\FunctionReturnTypeProviderInterface;
+use Psalm\Type;
+
+final class FunctionReturnTypeProvider implements FunctionReturnTypeProviderInterface
+{
+    /**
+     * @return array<lowercase-string>
+     */
+    public static function getFunctionIds(): array
+    {
+        return [
+            'psl\regex\capture_groups'
+        ];
+    }
+
+    public static function getFunctionReturnType(FunctionReturnTypeProviderEvent $event): ?Type\Union
+    {
+        $statements_source = $event->getStatementsSource();
+        $call_args = $event->getCallArgs();
+
+        $argument = $call_args[0] ?? null;
+        if (null === $argument) {
+            return self::fallbackType();
+        }
+
+        $argument_value = $argument->value;
+        $type = $statements_source->getNodeTypeProvider()->getType($argument_value);
+        if (null === $type) {
+            return self::fallbackType();
+        }
+
+        $atomic = $type->getAtomicTypes();
+        $capture_groups = $atomic['array'] ?? null;
+        if (!$capture_groups instanceof Type\Atomic\TKeyedArray) {
+            return self::fallbackType();
+        }
+
+        $string = static fn (): Type\Union => new Type\Union([new Type\Atomic\TString()]);
+        $properties = [
+            0 => $string()
+        ];
+        foreach ($capture_groups->properties as $value) {
+            $type = array_values($value->getAtomicTypes())[0] ?? null;
+            if (!$type instanceof Type\Atomic\TLiteralInt && !$type instanceof Type\Atomic\TLiteralString) {
+                return self::fallbackType();
+            }
+
+            $name = $type->value;
+
+            $properties[$name] = $string();
+        }
+
+        return new Type\Union([new Type\Atomic\TGenericObject('Psl\Type\TypeInterface', [
+            new Type\Union([
+                new Type\Atomic\TKeyedArray($properties)
+            ])
+        ])]);
+    }
+
+    private static function fallbackType(): Type\Union
+    {
+        return new Type\Union([new Type\Atomic\TGenericObject('Psl\Type\TypeInterface', [
+            new Type\Union([
+                new Type\Atomic\TArray([
+                    new Type\Union([new Type\Atomic\TArrayKey()]),
+                    new Type\Union([new Type\Atomic\TString()])
+                ])
+            ])
+        ])]);
+    }
+}
diff --git a/src/Plugin.php b/src/Plugin.php
index 6f8aa9f..3158a39 100644
--- a/src/Plugin.php
+++ b/src/Plugin.php
@@ -34,6 +34,9 @@ final class Plugin implements PluginEntryPointInterface
         yield EventHandler\Iter\Last\FunctionReturnTypeProvider::class;
         yield EventHandler\Iter\Count\FunctionReturnTypeProvider::class;
 
+        // Psl\Regex hooks
+        yield EventHandler\Regex\CaptureGroups\FunctionReturnTypeProvider::class;
+
         // Psl\Str hooks
         yield EventHandler\Str\After\FunctionReturnTypeProvider::class;
         yield EventHandler\Str\Before\FunctionReturnTypeProvider::class;