Component::lookup return type extension

Author: adaamz

README.md

@@ -12,6 +12,7 @@ This extension provides following features:
 * `Nette\ComponentModel\Container::getComponent()` knows type of the component because it reads the return type on `createComponent*` (this works best in presenters and controls)
 * `Nette\DI\Container::getByType` and `createInstance` return type based on first parameter (`Foo::class`).
 * `Nette\Forms\Container::getValues` return type based on `$asArray` parameter.
+* `Nette\ComponentModel\Component::lookup` return type based on `$throw` parameter.
 * Dynamic methods of [Nette\Utils\Html](https://doc.nette.org/en/2.4/html-elements)
 * Magic [Nette\Object and Nette\SmartObject](https://doc.nette.org/en/2.4/php-language-enhancements) properties
 * Event listeners through the `on*` properties

extension.neon

@@ -33,6 +33,11 @@ services:
 		tags:
 			- phpstan.broker.dynamicMethodReturnTypeExtension
 
+	-
+		class: PHPStan\Type\Nette\ComponentLookupDynamicReturnTypeExtension
+		tags:
+			- phpstan.broker.dynamicMethodReturnTypeExtension
+
 	-
 		class: PHPStan\Type\Nette\FormsBaseControlDynamicReturnTypeExtension
 		tags:

src/Type/Nette/ComponentLookupDynamicReturnTypeExtension.php

@@ -0,0 +1,46 @@
+<?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\Constant\ConstantBooleanType;
+use PHPStan\Type\DynamicMethodReturnTypeExtension;
+use PHPStan\Type\Type;
+use PHPStan\Type\TypeCombinator;
+
+final class ComponentLookupDynamicReturnTypeExtension implements DynamicMethodReturnTypeExtension
+{
+
+	public function getClass(): string
+	{
+		return \Nette\ComponentModel\Component::class;
+	}
+
+	public function isMethodSupported(MethodReflection $methodReflection): bool
+	{
+		return $methodReflection->getName() === 'lookup';
+	}
+
+	public function getTypeFromMethodCall(MethodReflection $methodReflection, MethodCall $methodCall, Scope $scope): Type
+	{
+		if (count($methodCall->args) < 2) {
+			return $methodReflection->getReturnType();
+		}
+
+		$paramNeedExpr = $methodCall->args[1]->value;
+		$paramNeedType = $scope->getType($paramNeedExpr);
+
+		if ($paramNeedType instanceof ConstantBooleanType) {
+			if ($paramNeedType->getValue()) {
+				return TypeCombinator::removeNull($methodReflection->getReturnType());
+			}
+
+			return TypeCombinator::addNull($methodReflection->getReturnType());
+		}
+
+		return $methodReflection->getReturnType();
+	}
+
+}