From d6aff3ed20598890c47255a85724fb763d679de7 Mon Sep 17 00:00:00 2001
From: shvlv <v.shelmuk@inpsyde.com>
Date: Tue, 24 Jan 2023 12:24:52 +0200
Subject: [PATCH] Parse class constant for PhpStorm Meta override

---
 .../Internal/Scanner/PhpStormMetaScanner.php  | 33 +++++++++++++++++++
 tests/StubTest.php                            | 16 +++++++++
 tests/fixtures/stubs/phpstorm.meta.php        |  7 ++++
 3 files changed, 56 insertions(+)

diff --git a/src/Psalm/Internal/Scanner/PhpStormMetaScanner.php b/src/Psalm/Internal/Scanner/PhpStormMetaScanner.php
index 5870f4925..e49ba2baf 100644
--- a/src/Psalm/Internal/Scanner/PhpStormMetaScanner.php
+++ b/src/Psalm/Internal/Scanner/PhpStormMetaScanner.php
@@ -12,6 +12,7 @@ use Psalm\Type\Atomic\TArray;
 use Psalm\Type\Atomic\TKeyedArray;
 use Psalm\Type\Atomic\TNamedObject;
 use Psalm\Type\Union;
+use ReflectionProperty;
 
 use function count;
 use function implode;
@@ -63,6 +64,38 @@ class PhpStormMetaScanner
                     } elseif ($array_item->value instanceof PhpParser\Node\Scalar\String_) {
                         $map[$array_item->key->value] = $array_item->value->value;
                     }
+                } elseif ($array_item->key instanceof PhpParser\Node\Expr\ClassConstFetch
+                    && $array_item->key->class instanceof PhpParser\Node\Name\FullyQualified
+                    && $array_item->key->name instanceof PhpParser\Node\Identifier
+                ) {
+                    $resolved_name =  $array_item->key->class->getAttribute('resolvedName');
+                    if (!$resolved_name) {
+                        continue;
+                    }
+
+                    $constant_type = $codebase->classlikes->getClassConstantType(
+                        $resolved_name,
+                        $array_item->key->name->name,
+                        ReflectionProperty::IS_PUBLIC,
+                    );
+
+                    if (!$constant_type instanceof Union || !$constant_type->isSingleStringLiteral()) {
+                        continue;
+                    }
+
+                    $meta_key = $constant_type->getSingleStringLiteral()->value;
+
+                    if ($array_item->value instanceof PhpParser\Node\Expr\ClassConstFetch
+                        && $array_item->value->class instanceof PhpParser\Node\Name\FullyQualified
+                        && $array_item->value->name instanceof PhpParser\Node\Identifier
+                        && strtolower($array_item->value->name->name)
+                    ) {
+                        $map[$meta_key] = new Union([
+                            new TNamedObject(implode('\\', $array_item->value->class->parts)),
+                        ]);
+                    } elseif ($array_item->value instanceof PhpParser\Node\Scalar\String_) {
+                        $map[$meta_key] = $array_item->value->value;
+                    }
                 }
             }
         }
diff --git a/tests/StubTest.php b/tests/StubTest.php
index 8cb13b69e..41ff56d2a 100644
--- a/tests/StubTest.php
+++ b/tests/StubTest.php
@@ -317,6 +317,10 @@ class StubTest extends TestCase
             '<?php
                 namespace Ns {
                     class MyClass {
+                    
+                        public const OBJECT = "object";
+                        public const EXCEPTION = "exception";
+                        
                         /**
                          * @return mixed
                          * @psalm-suppress InvalidReturnType
@@ -328,6 +332,12 @@ class StubTest extends TestCase
                          * @psalm-suppress InvalidReturnType
                          */
                         public function create2(string $s) {}
+                        
+                        /**
+                         * @return mixed
+                         * @psalm-suppress InvalidReturnType
+                         */
+                        public function create3(string $s) {}
 
                         /**
                          * @param mixed $s
@@ -374,6 +384,9 @@ class StubTest extends TestCase
 
                     $y1 = (new \Ns\MyClass)->creAte2("object");
                     $y2 = (new \Ns\MyClass)->creaTe2("exception");
+                    
+                    $const1 = (new \Ns\MyClass)->creAte3(\Ns\MyClass::OBJECT);
+                    $const2 = (new \Ns\MyClass)->creaTe3("exception");
 
                     $b1 = \Create("object");
                     $b2 = \cReate("exception");
@@ -404,6 +417,9 @@ class StubTest extends TestCase
                 '$y1===' => 'stdClass',
                 '$y2===' => 'Exception',
 
+                '$const1===' => 'stdClass',
+                '$const2===' => 'Exception',
+
                 '$b1===' => 'stdClass',
                 '$b2===' => 'Exception',
 
diff --git a/tests/fixtures/stubs/phpstorm.meta.php b/tests/fixtures/stubs/phpstorm.meta.php
index f2d3944f9..7ef519238 100644
--- a/tests/fixtures/stubs/phpstorm.meta.php
+++ b/tests/fixtures/stubs/phpstorm.meta.php
@@ -25,6 +25,13 @@ namespace PHPSTORM_META {
         'object' => \stdClass::class,
     ]));
 
+    // tests with class constant as key
+    override(\Ns\MyClass::crEate3(), map([
+        '' => '@',
+        \Ns\MyClass::EXCEPTION => \Exception::class,
+        \Ns\MyClass::OBJECT => \stdClass::class,
+    ]));
+
     override(\Ns\MyClass::foO(0), type(0));
     override(\Ns\MyClass::Bar(0), elementType(0));
     override(\foo(0), type(0));