Implement inline @var validation for return statements and add corresponding tests

Author: rcalicdan

src/Contract/DocblockExtractor.php

@@ -103,6 +103,29 @@ public static function extractTypeFromPropertyDoc(string $doc, string $propName)
         return null;
     }
 
+    /**
+     * Extracts the first @var tag's type string and variable name from a docblock string.
+     *
+     * @return array{0: string, 1: string}|null
+     */
+    public static function extractVarTagFromDoc(string $doc): ?array
+    {
+        try {
+            $phpDocNode = self::parseDocString($doc);
+            $varTags = $phpDocNode->getVarTagValues();
+            if (\count($varTags) > 0) {
+                $typeString = (string) $varTags[0]->type;
+                $varName = ltrim($varTags[0]->variableName, '$');
+
+                return [$typeString, $varName];
+            }
+        } catch (\Throwable $e) {
+            // Silently ignore malformed docblocks
+        }
+
+        return null;
+    }
+
     /**
      * Extracts local and imported type aliases (@phpstan-type and @phpstan-import-type) from a PHPDoc node.
      *
@@ -226,4 +249,4 @@ public static function extractMagicMethodContract(string $doc, string $methodNam
 
         return null;
     }
-}
+}

src/Internal/Checker/InlineChecker.php

@@ -98,18 +98,22 @@ public static function checkVariable(mixed $value, string $typeString, string $v
                 return $value;
             }
 
+            $context = ($varName === 'return') ? 'Return value' : "Variable \$$varName";
+
             if ($typeNode instanceof CallableTypeNode || ($typeNode instanceof IdentifierTypeNode && strtolower($typeNode->name) === 'callable')) {
-                return CallableWrapper::wrapTypeNode($typeNode, $value, "Variable \$$varName: Callback", $registry);
+                $cbPrefix = ($varName === 'return') ? 'Return value: Callback' : "Variable \$$varName: Callback";
+
+                return CallableWrapper::wrapTypeNode($typeNode, $value, $cbPrefix, $registry);
             }
 
             if ($typeNode instanceof GenericTypeNode && $checkGenerics && \is_object($value)) {
-                $err = TemplateManager::bindInstanceFromNode($value, $typeNode, "Variable \$$varName", true);
+                $err = TemplateManager::bindInstanceFromNode($value, $typeNode, $context, true);
                 if ($err !== null) {
                     return $err;
                 }
             }
 
-            $err = $registry->validate($value, $typeNode, "Variable \$$varName");
+            $err = $registry->validate($value, $typeNode, $context);
             if ($err !== null) {
                 return $err;
             }
@@ -305,4 +309,4 @@ private static function shouldValidateType(TypeNode $node, array $config): bool
 
         return (bool) ($config['scalars'] ?? false);
     }
-}
+}

src/Internal/ContractVisitor.php

@@ -7,6 +7,7 @@
 use PhpParser\Node;
 use PhpParser\NodeTraverser;
 use PhpParser\NodeVisitorAbstract;
+use TypePHP\Contract\DocblockExtractor;
 use TypePHP\Internal\Visitor\FunctionContractInjector;
 use TypePHP\Internal\Visitor\NodeBuilder;
 use TypePHP\Internal\Visitor\PropertyHookInjector;
@@ -68,6 +69,19 @@ public function enterNode(Node $node): array|int|null
             return null;
         }
 
+        if ($node instanceof Node\Stmt\Return_ && $node->expr !== null) {
+            $doc = $node->getDocComment();
+            if ($doc !== null && str_contains($doc->getText(), '@var')) {
+                $extracted = DocblockExtractor::extractVarTagFromDoc($doc->getText());
+                if ($extracted !== null) {
+                    [$typeString, $varName] = $extracted;
+                    $effectiveVarName = ($varName !== '') ? $varName : 'return';
+                    $checkCall = NodeBuilder::createVariableCheckCall($node->expr, $typeString, $effectiveVarName);
+                    $node->expr = NodeBuilder::createTernaryThrowExpr($checkCall, $node->getStartLine());
+                }
+            }
+        }
+
         if ($node instanceof Node\Stmt\Expression) {
             $doc = $node->getDocComment();
             if ($doc !== null && str_contains($doc->getText(), '@var')) {
@@ -214,4 +228,4 @@ private function extractDestructuringVariables(Node\Expr\List_|Node\Expr\Array_
 
         return $vars;
     }
-}
+}

tests/TypeChecking/Boundaries/InlineReturnValidationTest.php

@@ -0,0 +1,64 @@
+<?php
+
+declare(strict_types=1);
+
+/**
+ * Function with broad return type, but specific inline @var on return statement
+ */
+function testInlineVarOnReturnArray(): array
+{
+    /** @var list<string> */
+    return [1, 2, 3]; 
+}
+
+/**
+ * Valid inline @var on return statement
+ */
+function testValidInlineVarOnReturn(): array
+{
+    /** @var list<string> */
+    return ['apple', 'banana', 'cherry'];
+}
+
+/**
+ * Named @var tag on return statement
+ */
+function testNamedInlineVarOnReturn(int $id): int
+{
+    /** @var positive-int $id */
+    return $id;
+}
+
+/**
+ * Inline @var array shape on return inside a closure
+ */
+function testInlineVarOnReturnInClosure(): array
+{
+    $closure = function (): array {
+        /** @var array{id: positive-int, name: non-empty-string} */
+        return ['id' => -10, 'name' => 'Alice']; 
+    };
+
+    return $closure();
+}
+
+describe('Inline @var Validation on Direct Return Statements', function () {
+    test('validates and accepts valid return expression with inline @var', function () {
+        expect(testValidInlineVarOnReturn())->toBe(['apple', 'banana', 'cherry']);
+    });
+
+    test('throws TypeError when direct return expression violates unnamed inline @var contract', function () {
+        expect(fn () => testInlineVarOnReturnArray())
+            ->toThrow(TypeError::class, 'must be of type string');
+    });
+
+    test('throws TypeError when direct return expression violates named inline @var contract', function () {
+        expect(fn () => testNamedInlineVarOnReturn(-5))
+            ->toThrow(TypeError::class, 'positive-int');
+    });
+
+    test('throws TypeError when closure return expression violates inline @var array shape', function () {
+        expect(fn () => testInlineVarOnReturnInClosure())
+            ->toThrow(TypeError::class, 'positive-int');
+    });
+});