Add return type extension for constant()

Author: herndlm

conf/config.neon

@@ -1363,6 +1363,14 @@ services:
 		arguments:
 			checkMaybeUndefinedVariables: %checkMaybeUndefinedVariables%
 
+	-
+		class: PHPStan\Type\Php\ConstantFunctionReturnTypeExtension
+		tags:
+			- phpstan.broker.dynamicFunctionReturnTypeExtension
+
+	-
+	    class: PHPStan\Type\Php\ConstantHelper
+
 	-
 		class: PHPStan\Type\Php\CountFunctionReturnTypeExtension
 		tags:

src/Type/Php/ConstantFunctionReturnTypeExtension.php

@@ -0,0 +1,49 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Php;
+
+use PhpParser\Node\Expr\FuncCall;
+use PHPStan\Analyser\Scope;
+use PHPStan\Reflection\FunctionReflection;
+use PHPStan\Type\DynamicFunctionReturnTypeExtension;
+use PHPStan\Type\Type;
+use PHPStan\Type\TypeCombinator;
+use function count;
+
+class ConstantFunctionReturnTypeExtension implements DynamicFunctionReturnTypeExtension
+{
+
+	public function __construct(private ConstantHelper $constantHelper)
+	{
+	}
+
+	public function isFunctionSupported(FunctionReflection $functionReflection): bool
+	{
+		return $functionReflection->getName() === 'constant';
+	}
+
+	public function getTypeFromFunctionCall(
+		FunctionReflection $functionReflection,
+		FuncCall $functionCall,
+		Scope $scope,
+	): ?Type
+	{
+		if (count($functionCall->getArgs()) < 1) {
+			return null;
+		}
+
+		$nameType = $scope->getType($functionCall->getArgs()[0]->value);
+
+		$results = [];
+		foreach ($nameType->getConstantStrings() as $constantName) {
+			$results[] = $scope->getType($this->constantHelper->createExprFromConstantName($constantName->getValue()));
+		}
+
+		if (count($results) > 0) {
+			return TypeCombinator::union(...$results);
+		}
+
+		return null;
+	}
+
+}

src/Type/Php/ConstantHelper.php

@@ -0,0 +1,33 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Php;
+
+use PhpParser\Node\Expr;
+use PhpParser\Node\Expr\ClassConstFetch;
+use PhpParser\Node\Expr\ConstFetch;
+use PhpParser\Node\Identifier;
+use PhpParser\Node\Name;
+use PhpParser\Node\Name\FullyQualified;
+use function count;
+use function explode;
+use function ltrim;
+
+class ConstantHelper
+{
+
+	public function createExprFromConstantName(string $constantName): Expr
+	{
+		$classConstParts = explode('::', $constantName);
+		if (count($classConstParts) >= 2) {
+			$classConstName = new FullyQualified(ltrim($classConstParts[0], '\\'));
+			if ($classConstName->isSpecialClassName()) {
+				$classConstName = new Name($classConstName->toString());
+			}
+
+			return new ClassConstFetch($classConstName, new Identifier($classConstParts[1]));
+		}
+
+		return new ConstFetch(new FullyQualified($constantName));
+	}
+
+}

src/Type/Php/DefinedConstantTypeSpecifyingExtension.php

@@ -2,7 +2,6 @@
 
 namespace PHPStan\Type\Php;
 
-use PhpParser\Node;
 use PhpParser\Node\Expr\FuncCall;
 use PHPStan\Analyser\Scope;
 use PHPStan\Analyser\SpecifiedTypes;
@@ -14,14 +13,16 @@
 use PHPStan\Type\FunctionTypeSpecifyingExtension;
 use PHPStan\Type\MixedType;
 use function count;
-use function explode;
-use function ltrim;
 
 class DefinedConstantTypeSpecifyingExtension implements FunctionTypeSpecifyingExtension, TypeSpecifierAwareExtension
 {
 
 	private TypeSpecifier $typeSpecifier;
 
+	public function __construct(private ConstantHelper $constantHelper)
+	{
+	}
+
 	public function setTypeSpecifier(TypeSpecifier $typeSpecifier): void
 	{
 		$this->typeSpecifier = $typeSpecifier;
@@ -53,24 +54,8 @@ public function specifyTypes(
 			return new SpecifiedTypes([], []);
 		}
 
-		$classConstParts = explode('::', $constantName->getValue());
-		if (count($classConstParts) >= 2) {
-			$classConstName = new Node\Name\FullyQualified(ltrim($classConstParts[0], '\\'));
-			if ($classConstName->isSpecialClassName()) {
-				$classConstName = new Node\Name($classConstName->toString());
-			}
-			$constNode = new Node\Expr\ClassConstFetch(
-				$classConstName,
-				new Node\Identifier($classConstParts[1]),
-			);
-		} else {
-			$constNode = new Node\Expr\ConstFetch(
-				new Node\Name\FullyQualified($constantName->getValue()),
-			);
-		}
-
 		return $this->typeSpecifier->create(
-			$constNode,
+			$this->constantHelper->createExprFromConstantName($constantName->getValue()),
 			new MixedType(),
 			$context,
 			false,

tests/PHPStan/Analyser/NodeScopeResolverTest.php

@@ -642,6 +642,7 @@ public function dataFileAsserts(): iterable
 		yield from $this->gatherAssertTypes(__DIR__ . '/data/filter-var-array.php');
 
 		if (PHP_VERSION_ID >= 80100) {
+			yield from $this->gatherAssertTypes(__DIR__ . '/data/constant.php');
 			yield from $this->gatherAssertTypes(__DIR__ . '/data/enums.php');
 			yield from $this->gatherAssertTypes(__DIR__ . '/data/enums-import-alias.php');
 			yield from $this->gatherAssertTypes(__DIR__ . '/data/bug-7176.php');

tests/PHPStan/Analyser/data/constant.php

@@ -0,0 +1,40 @@
+<?php // lint >= 8.1
+
+namespace Constant;
+
+use function PHPStan\Testing\assertType;
+
+define('FOO', 'foo');
+const BAR = 'bar';
+
+class Baz
+{
+	const BAZ = 'baz';
+}
+
+enum Suit
+{
+    case Hearts;
+}
+
+function doFoo(string $constantName): void
+{
+	assertType('mixed', constant($constantName));
+}
+
+assertType("'foo'", FOO);
+assertType("'foo'", constant('FOO'));
+assertType("*ERROR*", constant('\Constant\FOO'));
+
+assertType("'bar'", BAR);
+assertType("*ERROR*", constant('BAR'));
+assertType("'bar'", constant('\Constant\BAR'));
+
+assertType("'bar'|'foo'", constant(rand(0, 1) ? 'FOO' : '\Constant\BAR'));
+
+assertType("'baz'", constant('\Constant\Baz::BAZ'));
+
+assertType('Constant\Suit::Hearts', Suit::Hearts);
+assertType('Constant\Suit::Hearts', constant('\Constant\Suit::Hearts'));
+
+assertType('*ERROR*', constant('UNDEFINED'));