From bd7f3424914807ad7f3c3355707b3b9bb4d7e76f Mon Sep 17 00:00:00 2001
From: Matthew Brown <github@muglug.com>
Date: Wed, 8 Jan 2020 23:48:42 -0500
Subject: [PATCH] Add constant definitions

---
 src/Psalm/Codebase.php                        |  8 +++
 .../Stubs/Generator/StubsGenerator.php        | 53 ++++++++++++++++---
 2 files changed, 54 insertions(+), 7 deletions(-)

diff --git a/src/Psalm/Codebase.php b/src/Psalm/Codebase.php
index 284e686b2..4261479dc 100644
--- a/src/Psalm/Codebase.php
+++ b/src/Psalm/Codebase.php
@@ -665,6 +665,14 @@ class Codebase
         return isset(self::$stubbed_constants[$const_id]) ? self::$stubbed_constants[$const_id] : null;
     }
 
+    /**
+     * @return array<string, Type\Union>
+     */
+    public function getAllStubbedConstants()
+    {
+        return self::$stubbed_constants;
+    }
+
     /**
      * @param  string $file_path
      *
diff --git a/src/Psalm/Internal/Stubs/Generator/StubsGenerator.php b/src/Psalm/Internal/Stubs/Generator/StubsGenerator.php
index c5bb514a9..287181624 100644
--- a/src/Psalm/Internal/Stubs/Generator/StubsGenerator.php
+++ b/src/Psalm/Internal/Stubs/Generator/StubsGenerator.php
@@ -14,13 +14,11 @@ class StubsGenerator
         \Psalm\Internal\Provider\ClassLikeStorageProvider $class_provider,
         \Psalm\Internal\Provider\FileStorageProvider $file_provider
     ) : string {
-        $all_class_storage = $class_provider->getAll();
-
         $namespaced_nodes = [];
 
         $psalm_base = dirname(__DIR__, 5);
 
-        foreach ($all_class_storage as $storage) {
+        foreach ($class_provider->getAll() as $storage) {
             if (\strpos($storage->name, 'Psalm\\') === 0) {
                 continue;
             }
@@ -51,12 +49,9 @@ class StubsGenerator
             );
         }
 
-        unset($all_class_storage);
-        $all_stubbed_functions = $codebase->functions->getAllStubbedFunctions();
-
         $all_function_names = [];
 
-        foreach ($all_stubbed_functions as $function_storage) {
+        foreach ($codebase->functions->getAllStubbedFunctions() as $function_storage) {
             if ($function_storage->location
                 && \strpos($function_storage->location->file_path, $psalm_base) === 0
             ) {
@@ -83,6 +78,26 @@ class StubsGenerator
             );
         }
 
+        foreach ($codebase->getAllStubbedConstants() as $fq_name => $type) {
+            if ($type->isMixed()) {
+                continue;
+            }
+
+            $name_parts = explode('\\', $fq_name);
+            $constant_name = array_pop($name_parts);
+
+            $namespace_name = implode('\\', $name_parts);
+
+            $namespaced_nodes[$namespace_name][$fq_name] = new PhpParser\Node\Stmt\Const_(
+                [
+                    new PhpParser\Node\Const_(
+                        $constant_name,
+                        self::getExpressionFromType($type)
+                    )
+                ]
+            );
+        }
+
         foreach ($file_provider->getAll() as $file_storage) {
             if (\strpos($file_storage->file_path, $psalm_base) === 0) {
                 continue;
@@ -112,6 +127,30 @@ class StubsGenerator
                     $namespace_name
                 );
             }
+
+            foreach ($file_storage->constants as $fq_name => $type) {
+                if ($type->isMixed()) {
+                    continue;
+                }
+
+                if ($type->isMixed()) {
+                    continue;
+                }
+
+                $name_parts = explode('\\', $fq_name);
+                $constant_name = array_pop($name_parts);
+
+                $namespace_name = implode('\\', $name_parts);
+
+                $namespaced_nodes[$namespace_name][$fq_name] = new PhpParser\Node\Stmt\Const_(
+                    [
+                        new PhpParser\Node\Const_(
+                            $constant_name,
+                            self::getExpressionFromType($type)
+                        )
+                    ]
+                );
+            }
         }
 
         ksort($namespaced_nodes);