Using messenger HandleTrait as QueryBus with appropriate result typing

Author: bnowak

README.md

@@ -29,6 +29,7 @@ This extension provides following features:
 * Provides correct return type for `Extension::getConfiguration()` method.
 * Provides correct return type for `CacheInterface::get()` method based on the callback return type.
 * Provides correct return type for `BrowserKitAssertionsTrait::getClient()` method.
+* Provides configurable return type resolution for methods that internally use Messenger `HandleTrait`.
 * Notifies you when you try to get an unregistered service from the container.
 * Notifies you when you try to get a private service from the container.
 * Notifies you when you access undefined console command arguments or options.
@@ -180,3 +181,95 @@ Call the new env in your `console-application.php`:
 ```php
 $kernel = new \App\Kernel('phpstan_env', (bool) $_SERVER['APP_DEBUG']);
 ```
+
+## Messenger HandleTrait Wrappers
+
+The extension provides advanced type inference for methods that internally use Symfony Messenger's `HandleTrait`. This feature is particularly useful for query bus implementations (in CQRS pattern) that use/wrap the `HandleTrait::handle()` method.
+
+### Configuration
+
+```neon
+parameters:
+    symfony:
+        messenger:
+            handleTraitWrappers:
+                - App\Bus\QueryBus::dispatch
+                - App\Bus\QueryBus::execute
+                - App\Bus\QueryBusInterface::dispatch
+```
+
+### Message Handlers
+
+```php
+use Symfony\Component\Messenger\Attribute\AsMessageHandler;
+
+// Product handler that returns Product
+#[AsMessageHandler]
+class GetProductQueryHandler
+{
+    public function __invoke(GetProductQuery $query): Product
+    {
+        return $this->productRepository->get($query->productId);
+    }
+}
+```
+
+### PHP Examples
+
+```php
+use Symfony\Component\Messenger\HandleTrait;
+use Symfony\Component\Messenger\MessageBusInterface;
+
+// Basic query bus implementation
+class QueryBus
+{
+    use HandleTrait;
+
+    public function __construct(MessageBusInterface $messageBus)
+    {
+        $this->messageBus = $messageBus;
+    }
+
+    public function dispatch(object $query): mixed
+    {
+        return $this->handle($query); // Return type will be inferred in calling code as query result
+    }
+
+    // Multiple methods per class example
+    public function execute(object $message): mixed
+    {
+        return $this->handle($message); // Return type will be inferred in calling code as query result
+    }
+}
+
+// Interface-based configuration example
+interface QueryBusInterface
+{
+    public function dispatch(object $query): mixed; // Return type will be inferred in calling code as query result
+}
+
+class QueryBusWithInterface implements QueryBusInterface
+{
+    use HandleTrait;
+
+    public function __construct(MessageBusInterface $queryBus)
+    {
+        $this->messageBus = $queryBus;
+    }
+
+    public function dispatch(object $query): mixed
+    {
+        return $this->handle($query);
+    }
+}
+
+// Examples of use with proper type inference
+$query = new GetProductQuery($productId);
+$queryBus = new QueryBus($messageBus);
+$queryBusWithInterface = new QueryBusWithInterface($messageBus);
+
+$product = $queryBus->dispatch($query); // Returns: Product
+$product2 = $queryBus->execute($query); // Returns: Product
+$product3 = $queryBusWithInterface->dispatch($query); // Returns: Product
+// Without the feature all above query bus results would be default 'mixed'.
+```

extension.neon

@@ -8,6 +8,8 @@ parameters:
 		containerXmlPath: null
 		constantHassers: true
 		consoleApplicationLoader: null
+		messenger:
+			handleTraitWrappers: []
 	stubFiles:
 		- stubs/Psr/Cache/CacheException.stub
 		- stubs/Psr/Cache/CacheItemInterface.stub
@@ -96,6 +98,9 @@ parametersSchema:
 		containerXmlPath: schema(string(), nullable())
 		constantHassers: bool()
 		consoleApplicationLoader: schema(string(), nullable())
+		messenger: structure([
+			handleTraitWrappers: listOf(string())
+		])
 	])
 
 services:
@@ -203,6 +208,13 @@ services:
 		class: PHPStan\Type\Symfony\MessengerHandleTraitReturnTypeExtension
 		tags: [phpstan.broker.expressionTypeResolverExtension]
 
+	# Messenger HandleTrait wrappers return type
+	-
+		class: PHPStan\Type\Symfony\MessengerHandleTraitWrapperReturnTypeExtension
+		tags: [phpstan.broker.expressionTypeResolverExtension]
+		arguments:
+			messenger: %symfony.messenger%
+
 	# InputInterface::getArgument() return type
 	-
 		factory: PHPStan\Type\Symfony\InputInterfaceGetArgumentDynamicReturnTypeExtension

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;
 
@@ -30,26 +31,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

@@ -0,0 +1,144 @@
+<?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\Reflection\ReflectionProvider;
+use PHPStan\Symfony\MessageMap;
+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;
+
+/**
+ * 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
+{
+
+	private MessageMapFactory $messageMapFactory;
+
+	private ?MessageMap $messageMap = null;
+
+	/** @var array<string> */
+	private array $wrappers;
+
+	private ReflectionProvider $reflectionProvider;
+
+	/** @param array{handleTraitWrappers: array<string>}|null $messenger */
+	public function __construct(MessageMapFactory $messageMapFactory, ?array $messenger, ReflectionProvider $reflectionProvider)
+	{
+		$this->messageMapFactory = $messageMapFactory;
+		$this->wrappers = $messenger['handleTraitWrappers'] ?? [];
+		$this->reflectionProvider = $reflectionProvider;
+	}
+
+	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) === 0) {
+			return null;
+		}
+
+		$returnTypes = [];
+		foreach ($argClassNames as $argClassName) {
+			$messageMap = $this->getMessageMap();
+			$returnType = $messageMap->getTypeForClass($argClassName);
+
+			if (is_null($returnType)) {
+				return null;
+			}
+
+			$returnTypes[] = $returnType;
+		}
+
+		return TypeCombinator::union(...$returnTypes);
+	}
+
+	/**
+	 * @phpstan-assert-if-true =MethodCall $expr
+	 */
+	private function isSupported(Expr $expr, Scope $scope): bool
+	{
+		if ($this->wrappers === []) {
+			return false;
+		}
+
+		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) === 0) {
+			return false;
+		}
+
+		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
+		if (in_array($classMethodCombination, $this->wrappers, true)) {
+			return true;
+		}
+
+		// Check if any interface implemented by this class::method is configured
+		if ($this->reflectionProvider->hasClass($className)) {
+			$classReflection = $this->reflectionProvider->getClass($className);
+			foreach ($classReflection->getInterfaces() as $interface) {
+				$interfaceMethodCombination = $interface->getName() . '::' . $methodName;
+				if (in_array($interfaceMethodCombination, $this->wrappers, true)) {
+					return true;
+				}
+			}
+		}
+
+		return false;
+	}
+
+	private function getMessageMap(): MessageMap
+	{
+		if ($this->messageMap === null) {
+			$this->messageMap = $this->messageMapFactory->create();
+		}
+
+		return $this->messageMap;
+	}
+
+}