Dynamic return type extension for DI Container getByType/createInstance methods

Author: ondrejmirtes

composer.json

@@ -8,6 +8,7 @@
 		"php": "~7.0",
 		"phpstan/phpstan": "^0.6",
 		"nette/component-model": "^2.3.0",
+		"nette/di": "^2.3.0",
 		"nette/forms": "^2.3.0",
 		"nette/utils": "^2.3.0"
 	},

extension.neon

@@ -33,3 +33,8 @@ services:
 		class: PHPStan\Type\Nette\FormsBaseControlDynamicReturnTypeExtension
 		tags:
 			- phpstan.broker.dynamicMethodReturnTypeExtension
+
+	-
+		class: PHPStan\Type\Nette\ServiceLocatorDynamicReturnTypeExtension
+		tags:
+			- phpstan.broker.dynamicMethodReturnTypeExtension

src/Type/Nette/ServiceLocatorDynamicReturnTypeExtension.php

@@ -0,0 +1,65 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Nette;
+
+use PhpParser\Node\Expr\MethodCall;
+use PHPStan\Analyser\Scope;
+use PHPStan\Reflection\MethodReflection;
+use PHPStan\Type\ObjectType;
+use PHPStan\Type\TrueBooleanType;
+use PHPStan\Type\Type;
+
+class ServiceLocatorDynamicReturnTypeExtension implements \PHPStan\Type\DynamicMethodReturnTypeExtension
+{
+
+	public static function getClass(): string
+	{
+		return \Nette\DI\Container::class;
+	}
+
+	public function isMethodSupported(MethodReflection $methodReflection): bool
+	{
+		return in_array($methodReflection->getName(), [
+			'getByType',
+			'createInstance',
+		], true);
+	}
+
+	public function getTypeFromMethodCall(MethodReflection $methodReflection, MethodCall $methodCall, Scope $scope): Type
+	{
+		if (count($methodCall->args) === 0) {
+			return $methodReflection->getReturnType();
+		}
+		$arg = $methodCall->args[0]->value;
+		if (!($arg instanceof \PhpParser\Node\Expr\ClassConstFetch)) {
+			return $methodReflection->getReturnType();
+		}
+
+		$class = $arg->class;
+		if (!($class instanceof \PhpParser\Node\Name)) {
+			return $methodReflection->getReturnType();
+		}
+
+		$class = (string) $class;
+
+		if ($class === 'static') {
+			return $methodReflection->getReturnType();
+		}
+
+		if ($class === 'self') {
+			$class = $scope->getClass();
+		}
+
+		$isNullable = false;
+		if (
+			$methodReflection->getName() === 'getByType'
+			&& count($methodCall->args) >= 2
+			&& $scope->getType($methodCall->args[1]->value) instanceof TrueBooleanType
+		) {
+			$isNullable = true;
+		}
+
+		return new ObjectType($class, $isNullable);
+	}
+
+}