From bd8fca805d11411df1d566ff6fb8be7de9647559 Mon Sep 17 00:00:00 2001
From: Andres Kalle <andres@play.ee>
Date: Thu, 14 May 2026 14:59:32 +0300
Subject: [PATCH] Add dynamic return type extension for Psl\Type\union()

---
 extension.neon                                |  5 ++
 src/Type/TypeUnionReturnTypeExtension.php     | 48 +++++++++++++++++++
 tests/Type/PslTypeSpecifyingExtensionTest.php |  1 +
 tests/Type/data/unionV2.php                   | 47 ++++++++++++++++++
 4 files changed, 101 insertions(+)
 create mode 100644 src/Type/TypeUnionReturnTypeExtension.php
 create mode 100644 tests/Type/data/unionV2.php

diff --git a/extension.neon b/extension.neon
index 8f3a510..b2ad1a8 100644
--- a/extension.neon
+++ b/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:
diff --git a/src/Type/TypeUnionReturnTypeExtension.php b/src/Type/TypeUnionReturnTypeExtension.php
new file mode 100644
index 0000000..ec2e1de
--- /dev/null
+++ b/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),
+			]
+		);
+	}
+
+}
diff --git a/tests/Type/PslTypeSpecifyingExtensionTest.php b/tests/Type/PslTypeSpecifyingExtensionTest.php
index 3944237..6acbbfc 100644
--- a/tests/Type/PslTypeSpecifyingExtensionTest.php
+++ b/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');
 		}
 	}
 
diff --git a/tests/Type/data/unionV2.php b/tests/Type/data/unionV2.php
new file mode 100644
index 0000000..bb0598b
--- /dev/null
+++ b/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);
+	}
+}