implementation

Author: Bartłomiej Nowak

extension.neon

@@ -11,6 +11,11 @@ parameters:
 		constantHassers: true
 		console_application_loader: null
 		consoleApplicationLoader: null
+		messenger:
+			handleTraitWrappers:
+			    # move that params to tests only
+				- MessengerHandleTrait\QueryBus::dispatch
+				- MessengerHandleTrait\QueryBus2::dispatch
 	featureToggles:
 		skipCheckGenericClasses:
 			- Symfony\Component\Form\AbstractType
@@ -115,6 +120,9 @@ parametersSchema:
 		constantHassers: bool()
 		console_application_loader: schema(string(), nullable())
 		consoleApplicationLoader: schema(string(), nullable())
+		messenger: structure([
+			handleTraitWrappers: listOf(string())
+		])
 	])
 
 services:
@@ -214,6 +222,11 @@ services:
 		class: PHPStan\Type\Symfony\MessengerHandleTraitReturnTypeExtension
 		tags: [phpstan.broker.expressionTypeResolverExtension]
 
+	# Messenger HandleTrait wrappers return type
+	-
+		class: PHPStan\Type\Symfony\MessengerHandleTraitWrapperReturnTypeExtension
+		tags: [phpstan.broker.expressionTypeResolverExtension]
+
 	# InputInterface::getArgument() return type
 	-
 		factory: PHPStan\Type\Symfony\InputInterfaceGetArgumentDynamicReturnTypeExtension

src/Symfony/Configuration.php

@@ -31,4 +31,12 @@ public function getConsoleApplicationLoader(): ?string
 		return $this->parameters['consoleApplicationLoader'] ?? $this->parameters['console_application_loader'] ?? null;
 	}
 
+	/**
+	 * @return array<string>
+	 */
+	public function getMessengerHandleTraitWrappers(): array
+	{
+		return $this->parameters['messenger']['handleTraitWrappers'] ?? [];
+	}
+
 }

src/Type/Symfony/MessengerHandleTraitWrapperReturnTypeExtension.php

@@ -0,0 +1,105 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Symfony;
+
+use PhpParser\Node\Expr;
+use PhpParser\Node\Expr\MethodCall;
+use PhpParser\Node\Identifier;
+use PHPStan\Analyser\Scope;
+use PHPStan\Symfony\Configuration;
+use PHPStan\Symfony\MessageMap;
+use PHPStan\Symfony\MessageMapFactory;
+use PHPStan\Type\ExpressionTypeResolverExtension;
+use PHPStan\Type\Type;
+use function count;
+use function in_array;
+use function is_null;
+
+/**
+ * Configurable extension for resolving return types of methods that internally use HandleTrait.
+ *
+ * Configured via PHPStan parameters under symfony.messenger.handleTraitWrappers with
+ * "Class::method" patterns, e.g.:
+ * - App\Bus\QueryBus::dispatch
+ * - App\Bus\QueryBus::query
+ * - App\Bus\CommandBus::execute
+ * - App\Bus\CommandBus::handle
+ */
+final class MessengerHandleTraitWrapperReturnTypeExtension implements ExpressionTypeResolverExtension
+{
+
+	/** @var MessageMapFactory */
+	private $messageMapFactory;
+
+	/** @var MessageMap|null */
+	private $messageMap;
+
+	/** @var array<string> */
+	private $wrappers;
+
+	public function __construct(MessageMapFactory $messageMapFactory, Configuration $configuration)
+	{
+		$this->messageMapFactory = $messageMapFactory;
+		$this->wrappers = $configuration->getMessengerHandleTraitWrappers();
+	}
+
+	public function getType(Expr $expr, Scope $scope): ?Type
+	{
+		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();
+
+		if (count($argClassNames) === 1) {
+			$messageMap = $this->getMessageMap();
+			$returnType = $messageMap->getTypeForClass($argClassNames[0]);
+
+			if (!is_null($returnType)) {
+				return $returnType;
+			}
+		}
+
+		return null;
+	}
+
+	/**
+	 * @phpstan-assert-if-true =MethodCall $expr
+	 */
+	private function isSupported(Expr $expr, Scope $scope): bool
+	{
+		if (!($expr instanceof MethodCall) || !($expr->name instanceof Identifier)) {
+			return false;
+		}
+
+		$methodName = $expr->name->name;
+		$varType = $scope->getType($expr->var);
+		$classNames = $varType->getObjectClassNames();
+
+		if (count($classNames) !== 1) {
+			return false;
+		}
+
+		$className = $classNames[0];
+		$classMethodCombination = $className . '::' . $methodName;
+
+		// Check if this class::method combination is configured
+		return in_array($classMethodCombination, $this->wrappers, true);
+	}
+
+	private function getMessageMap(): MessageMap
+	{
+		if ($this->messageMap === null) {
+			$this->messageMap = $this->messageMapFactory->create();
+		}
+
+		return $this->messageMap;
+	}
+
+}

tests/Type/Symfony/data/messenger_handle_trait.php

@@ -121,6 +121,15 @@ public function dispatch(object $query): mixed
     }
 }
 
+class QueryBus2 {
+	use HandleTrait;
+
+	public function dispatch(object $query): mixed
+	{
+		return $this->handle($query);
+	}
+}
+
 class Controller {
     public function action()
     {
@@ -138,5 +147,8 @@ public function action()
         // HandleTrait will throw exception in fact due to multiple handle methods/handlers per single query
         assertType('mixed', $queryBus->dispatch(new MultiHandlesForInTheSameHandlerQuery()));
         assertType('mixed', $queryBus->dispatch(new MultiHandlersForTheSameMessageQuery()));
+
+		$queryBus2 = new QueryBus2();
+		assertType(TaggedResult::class, $queryBus2->dispatch(new TaggedQuery()));
     }
 }