Narrow Psl\Type\union() return type

extension.neon

@@ -14,6 +14,11 @@ services:
 		tags:
 			- phpstan.broker.dynamicFunctionReturnTypeExtension
 
+	-
+		class: Psl\PHPStan\Type\TypeUnionReturnTypeExtension
+		tags:
+			- phpstan.broker.dynamicFunctionReturnTypeExtension
+
 	-
 		class: Psl\PHPStan\Type\AssertTypeSpecifyingExtension
 		tags:

src/Type/TypeUnionReturnTypeExtension.php

@@ -0,0 +1,48 @@
+<?php declare(strict_types = 1);
+
+namespace Psl\PHPStan\Type;
+
+use PhpParser\Node\Expr\FuncCall;
+use PHPStan\Analyser\Scope;
+use PHPStan\Reflection\FunctionReflection;
+use PHPStan\Type\DynamicFunctionReturnTypeExtension;
+use PHPStan\Type\ErrorType;
+use PHPStan\Type\Generic\GenericObjectType;
+use PHPStan\Type\Type;
+use PHPStan\Type\TypeCombinator;
+use Psl\Type\TypeInterface;
+
+class TypeUnionReturnTypeExtension implements DynamicFunctionReturnTypeExtension
+{
+
+	public function isFunctionSupported(FunctionReflection $functionReflection): bool
+	{
+		return $functionReflection->getName() === 'Psl\Type\union';
+	}
+
+	public function getTypeFromFunctionCall(FunctionReflection $functionReflection, FuncCall $functionCall, Scope $scope): ?Type
+	{
+		$args = $functionCall->getArgs();
+		if ($args === []) {
+			return null;
+		}
+
+		$innerTypes = [];
+		foreach ($args as $arg) {
+			$argType = $scope->getType($arg->value);
+			$inner = $argType->getTemplateType(TypeInterface::class, 'T');
+			if ($inner instanceof ErrorType) {
+				return null;
+			}
+			$innerTypes[] = $inner;
+		}
+
+		return new GenericObjectType(
+			TypeInterface::class,
+			[
+				TypeCombinator::union(...$innerTypes),
+			]
+		);
+	}
+
+}

tests/Type/PslTypeSpecifyingExtensionTest.php

@@ -28,6 +28,7 @@ public function dataFileAsserts(): iterable
 			yield from $this->gatherAssertTypes(__DIR__ . '/data/complexTypev1.php');
 		} else {
 			yield from $this->gatherAssertTypes(__DIR__ . '/data/complexTypev2.php');
+			yield from $this->gatherAssertTypes(__DIR__ . '/data/unionV2.php');
 		}
 	}
 

tests/Type/data/unionV2.php

@@ -0,0 +1,47 @@
+<?php declare(strict_types=1);
+
+namespace PslUnionV2Test;
+
+use Psl\Type;
+
+use function PHPStan\Testing\assertType;
+
+/**
+ * For PSL >= 2.0.0 (literal_scalar)
+ */
+class UnionTypes
+{
+	public function literalScalarUnion($input): void
+	{
+		$ab = Type\union(Type\literal_scalar('a'), Type\literal_scalar('b'));
+		$out = $ab->coerce($input);
+		assertType("'a'|'b'", $out);
+	}
+
+	public function mixedUnion($input): void
+	{
+		$intOrString = Type\union(Type\int(), Type\string());
+		$out = $intOrString->coerce($input);
+		assertType('int|string', $out);
+	}
+
+	public function shapeWithLiteralUnion($input): void
+	{
+		$shape = Type\shape([
+			'kind' => Type\union(
+				Type\literal_scalar('a'),
+				Type\literal_scalar('b'),
+				Type\literal_scalar('c'),
+			),
+		]);
+		$out = $shape->coerce($input);
+		assertType("array{kind: 'a'|'b'|'c'}", $out);
+	}
+
+	public function singleArgUnion($input): void
+	{
+		$single = Type\union(Type\int());
+		$out = $single->coerce($input);
+		assertType('int', $out);
+	}
+}