Enum support in query type inference

Author: arnaud-lb

composer.json

@@ -27,6 +27,7 @@
 		"doctrine/persistence": "^1.3.8 || ^2.2.1",
 		"nesbot/carbon": "^2.49",
 		"nikic/php-parser": "^4.13.2",
+		"ocramius/package-versions": "*",
 		"php-parallel-lint/php-parallel-lint": "^1.2",
 		"phpstan/phpstan-phpunit": "^1.0",
 		"phpstan/phpstan-strict-rules": "^1.0",
@@ -41,7 +42,8 @@
 		},
 		"sort-packages": true,
 		"allow-plugins": {
-			"composer/package-versions-deprecated": true
+			"composer/package-versions-deprecated": true,
+			"ocramius/package-versions": true
 		}
 	},
 	"extra": {

phpstan.neon

@@ -17,6 +17,10 @@ parameters:
 
 	reportUnmatchedIgnoredErrors: false
 
+	bootstrapFiles:
+		- stubs/runtime/Enum/UnitEnum.php
+		- stubs/runtime/Enum/BackedEnum.php
+
 	ignoreErrors:
 		-
 			message: '~^Variable method call on Doctrine\\ORM\\QueryBuilder~'

src/Type/Doctrine/Query/QueryResultTypeWalker.php

@@ -2,8 +2,10 @@
 
 namespace PHPStan\Type\Doctrine\Query;
 
+use BackedEnum;
 use Doctrine\ORM\EntityManagerInterface;
 use Doctrine\ORM\Mapping\ClassMetadata;
+use Doctrine\ORM\Mapping\ClassMetadataInfo;
 use Doctrine\ORM\Query;
 use Doctrine\ORM\Query\AST;
 use Doctrine\ORM\Query\AST\TypedExpression;
@@ -15,6 +17,7 @@
 use PHPStan\Type\Constant\ConstantFloatType;
 use PHPStan\Type\Constant\ConstantIntegerType;
 use PHPStan\Type\Constant\ConstantStringType;
+use PHPStan\Type\ConstantTypeHelper;
 use PHPStan\Type\Doctrine\DescriptorNotRegisteredException;
 use PHPStan\Type\Doctrine\DescriptorRegistry;
 use PHPStan\Type\FloatType;
@@ -31,6 +34,7 @@
 use PHPStan\Type\TypeTraverser;
 use PHPStan\Type\TypeUtils;
 use PHPStan\Type\UnionType;
+use function array_map;
 use function assert;
 use function class_exists;
 use function count;
@@ -42,6 +46,7 @@
 use function is_numeric;
 use function is_object;
 use function is_string;
+use function is_subclass_of;
 use function serialize;
 use function sprintf;
 use function strtolower;
@@ -231,15 +236,13 @@ public function walkPathExpression($pathExpr)
 
 		switch ($pathExpr->type) {
 			case AST\PathExpression::TYPE_STATE_FIELD:
-				$typeName = $class->getTypeOfField($fieldName);
-
-				assert(is_string($typeName));
+				[$typeName, $enumType] = $this->getTypeOfField($class, $fieldName);
 
 				$nullable = $this->isQueryComponentNullable($dqlAlias)
 					|| $class->isNullable($fieldName)
 					|| $this->hasAggregateWithoutGroupBy();
 
-				$fieldType = $this->resolveDatabaseInternalType($typeName, $nullable);
+				$fieldType = $this->resolveDatabaseInternalType($typeName, $enumType, $nullable);
 
 				return $this->marshalType($fieldType);
 
@@ -273,14 +276,12 @@ public function walkPathExpression($pathExpr)
 				}
 
 				$targetFieldName = $identifierFieldNames[0];
-				$typeName = $targetClass->getTypeOfField($targetFieldName);
-
-				assert(is_string($typeName));
+				[$typeName] = $this->getTypeOfField($targetClass, $targetFieldName);
 
 				$nullable = (bool) ($joinColumn['nullable'] ?? true)
 					|| $this->hasAggregateWithoutGroupBy();
 
-				$fieldType = $this->resolveDatabaseInternalType($typeName, $nullable);
+				$fieldType = $this->resolveDatabaseInternalType($typeName, null, $nullable);
 
 				return $this->marshalType($fieldType);
 
@@ -556,7 +557,7 @@ public function walkFunction($function)
 				$nullable = (bool) ($joinColumn['nullable'] ?? true)
 					|| $this->hasAggregateWithoutGroupBy();
 
-				$fieldType = $this->resolveDatabaseInternalType($typeName, $nullable);
+				$fieldType = $this->resolveDatabaseInternalType($typeName, null, $nullable);
 
 				return $this->marshalType($fieldType);
 
@@ -783,15 +784,13 @@ public function walkSelectExpression($selectExpression)
 			$qComp = $this->queryComponents[$dqlAlias];
 			$class = $qComp['metadata'];
 
-			$typeName = $class->getTypeOfField($fieldName);
-
-			assert(is_string($typeName));
+			[$typeName, $enumType] = $this->getTypeOfField($class, $fieldName);
 
 			$nullable = $this->isQueryComponentNullable($dqlAlias)
 				|| $class->isNullable($fieldName)
 				|| $this->hasAggregateWithoutGroupBy();
 
-			$type = $this->resolveDoctrineType($typeName, $nullable);
+			$type = $this->resolveDoctrineType($typeName, $enumType, $nullable);
 
 			$this->typeBuilder->addScalar($resultAlias, $type);
 
@@ -1295,14 +1294,37 @@ private function isQueryComponentNullable(string $dqlAlias): bool
 		return $this->nullableQueryComponents[$dqlAlias] ?? false;
 	}
 
-	private function resolveDoctrineType(string $typeName, bool $nullable = false): Type
+	/** @return array{string, ?class-string<BackedEnum>} Doctrine type name and enum type of field */
+	private function getTypeOfField(ClassMetadataInfo $class, string $fieldName): array
 	{
-		try {
-			$type = $this->descriptorRegistry
-				->get($typeName)
-				->getWritableToPropertyType();
-		} catch (DescriptorNotRegisteredException $e) {
-			$type = new MixedType();
+		assert(isset($class->fieldMappings[$fieldName]));
+
+		/** @var array{type: string, enumType?: ?string} $metadata */
+		$metadata = $class->fieldMappings[$fieldName];
+
+		$type = $metadata['type'];
+		$enumType = $metadata['enumType'] ?? null;
+
+		if (!is_string($enumType) || !class_exists($enumType) || !is_subclass_of($enumType, BackedEnum::class)) {
+			$enumType = null;
+		}
+
+		return [$type, $enumType];
+	}
+
+	/** @param ?class-string<BackedEnum> $enumType */
+	private function resolveDoctrineType(string $typeName, ?string $enumType = null, bool $nullable = false): Type
+	{
+		if ($enumType !== null) {
+			$type = new ObjectType($enumType);
+		} else {
+			try {
+				$type = $this->descriptorRegistry
+					->get($typeName)
+					->getWritableToPropertyType();
+			} catch (DescriptorNotRegisteredException $e) {
+				$type = new MixedType();
+			}
 		}
 
 		if ($nullable) {
@@ -1312,7 +1334,8 @@ private function resolveDoctrineType(string $typeName, bool $nullable = false):
 		return $type;
 	}
 
-	private function resolveDatabaseInternalType(string $typeName, bool $nullable = false): Type
+	/** @param ?class-string<BackedEnum> $enumType */
+	private function resolveDatabaseInternalType(string $typeName, ?string $enumType = null, bool $nullable = false): Type
 	{
 		try {
 			$type = $this->descriptorRegistry
@@ -1322,6 +1345,15 @@ private function resolveDatabaseInternalType(string $typeName, bool $nullable =
 			$type = new MixedType();
 		}
 
+		if ($enumType !== null) {
+			$enumTypes = array_map(static function ($enumType) {
+				return ConstantTypeHelper::getTypeFromValue($enumType->value);
+			}, $enumType::cases());
+			$enumType = TypeCombinator::union(...$enumTypes);
+			$enumType = TypeCombinator::union($enumType, $enumType->toString());
+			$type = TypeCombinator::intersect($enumType, $type);
+		}
+
 		if ($nullable) {
 			$type = TypeCombinator::addNull($type);
 		}

stubs/runtime/Enum/BackedEnum.php

@@ -0,0 +1,22 @@
+<?php
+
+if (\PHP_VERSION_ID < 80100) {
+	if (interface_exists('BackedEnum', false)) {
+		return;
+	}
+
+	interface BackedEnum extends UnitEnum
+	{
+		/**
+		 * @param int|string $value
+		 * @return static
+		 */
+		public static function from($value);
+
+		/**
+		 * @param int|string $value
+		 * @return ?static
+		 */
+		public static function tryFrom($value);
+	}
+}

stubs/runtime/Enum/UnitEnum.php

@@ -0,0 +1,15 @@
+<?php
+
+if (\PHP_VERSION_ID < 80100) {
+	if (interface_exists('UnitEnum', false)) {
+		return;
+	}
+
+	interface UnitEnum
+	{
+		/**
+		 * @return static[]
+		 */
+		public static function cases(): array;
+	}
+}