Dynamic return type extension for fluent interfaces of form fields

Author: ondrejmirtes

composer.json

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

extension.neon

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

src/Type/Nette/FormsBaseControlDynamicReturnTypeExtension.php

@@ -0,0 +1,38 @@
+<?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\Type;
+
+class FormsBaseControlDynamicReturnTypeExtension implements \PHPStan\Type\DynamicMethodReturnTypeExtension
+{
+
+	public static function getClass(): string
+	{
+		return \Nette\Forms\Controls\BaseControl::class;
+	}
+
+	public function isMethodSupported(
+		MethodReflection $methodReflection
+	): bool
+	{
+		return $methodReflection->getDeclaringClass()->getName() === \Nette\Forms\Controls\BaseControl::class;
+	}
+
+	public function getTypeFromMethodCall(
+		MethodReflection $methodReflection,
+		MethodCall $methodCall,
+		Scope $scope
+	): Type
+	{
+		if ($methodReflection->getReturnType()->getClass() !== null && $methodReflection->getReturnType()->getClass() === \Nette\Forms\Controls\BaseControl::class) {
+			return $scope->getType($methodCall->var);
+		}
+
+		return $methodReflection->getReturnType();
+	}
+
+}