support for handling union types

Author: bnowak

src/Type/Symfony/MessengerHandleTraitReturnTypeExtension.php

@@ -10,6 +10,7 @@
 use PHPStan\Symfony\MessageMapFactory;
 use PHPStan\Type\ExpressionTypeResolverExtension;
 use PHPStan\Type\Type;
+use PHPStan\Type\TypeCombinator;
 use function count;
 use function is_null;
 
@@ -32,26 +33,36 @@ public function __construct(MessageMapFactory $symfonyMessageMapFactory)
 
 	public function getType(Expr $expr, Scope $scope): ?Type
 	{
-		if ($this->isSupported($expr, $scope)) {
-			$args = $expr->getArgs();
-			if (count($args) !== 1) {
-				return null;
-			}
+		if (!$this->isSupported($expr, $scope)) {
+			return null;
+		}
+
+		$args = $expr->getArgs();
+		if (count($args) !== 1) {
+			return null;
+		}
 
-			$arg = $args[0]->value;
-			$argClassNames = $scope->getType($arg)->getObjectClassNames();
+		$arg = $args[0]->value;
+		$argClassNames = $scope->getType($arg)->getObjectClassNames();
 
-			if (count($argClassNames) === 1) {
-				$messageMap = $this->getMessageMap();
-				$returnType = $messageMap->getTypeForClass($argClassNames[0]);
+		if (count($argClassNames) === 0) {
+			return null;
+		}
+
+		$messageMap = $this->getMessageMap();
+
+		$returnTypes = [];
+		foreach ($argClassNames as $argClassName) {
+			$returnType = $messageMap->getTypeForClass($argClassName);
 
-				if (!is_null($returnType)) {
-					return $returnType;
-				}
+			if (is_null($returnType)) {
+				return null;
 			}
+
+			$returnTypes[] = $returnType;
 		}
 
-		return null;
+		return TypeCombinator::union(...$returnTypes);
 	}
 
 	private function getMessageMap(): MessageMap

src/Type/Symfony/MessengerHandleTraitWrapperReturnTypeExtension.php

@@ -12,6 +12,7 @@
 use PHPStan\Symfony\MessageMapFactory;
 use PHPStan\Type\ExpressionTypeResolverExtension;
 use PHPStan\Type\Type;
+use PHPStan\Type\TypeCombinator;
 use function count;
 use function in_array;
 use function is_null;
@@ -62,16 +63,23 @@ public function getType(Expr $expr, Scope $scope): ?Type
 		$arg = $args[0]->value;
 		$argClassNames = $scope->getType($arg)->getObjectClassNames();
 
-		if (count($argClassNames) === 1) {
+		if (count($argClassNames) === 0) {
+			return null;
+		}
+
+		$returnTypes = [];
+		foreach ($argClassNames as $argClassName) {
 			$messageMap = $this->getMessageMap();
-			$returnType = $messageMap->getTypeForClass($argClassNames[0]);
+			$returnType = $messageMap->getTypeForClass($argClassName);
 
-			if (!is_null($returnType)) {
-				return $returnType;
+			if (is_null($returnType)) {
+				return null;
 			}
+
+			$returnTypes[] = $returnType;
 		}
 
-		return null;
+		return TypeCombinator::union(...$returnTypes);
 	}
 
 	/**
@@ -87,11 +95,21 @@ private function isSupported(Expr $expr, Scope $scope): bool
 		$varType = $scope->getType($expr->var);
 		$classNames = $varType->getObjectClassNames();
 
-		if (count($classNames) !== 1) {
+		if (count($classNames) === 0) {
 			return false;
 		}
 
-		$className = $classNames[0];
+		foreach ($classNames as $className) {
+			if (!$this->isClassMethodSupported($className, $methodName)) {
+				return false;
+			}
+		}
+
+		return true;
+	}
+
+	private function isClassMethodSupported(string $className, string $methodName): bool
+	{
 		$classMethodCombination = $className . '::' . $methodName;
 
 		// Check if this exact class::method combination is configured