Rector rule to add the `@return FormInterface<TData>` on twig component with form trait on method `instantiateForm`

src\Rector\AddFormInterfaceGenericToTwigComponentWithFormRector.php

<?php

declare(strict_types=1);

namespace Utils\Rector\Rector;

use PhpParser\Comment\Doc;
use PhpParser\Node;
use PhpParser\Node\Stmt\Class_;
use PhpParser\Node\Stmt\TraitUse;
use PhpParser\Node\Stmt\ClassMethod;
use PhpParser\Node\Expr\ClassConstFetch;
use PhpParser\Node\Expr\MethodCall;
use PHPStan\PhpDocParser\Ast\PhpDoc\ExtendsTagValueNode;
use PHPStan\PhpDocParser\Ast\PhpDoc\PhpDocTagNode;
use PHPStan\PhpDocParser\Ast\PhpDoc\ReturnTagValueNode;
use PHPStan\PhpDocParser\Ast\Type\GenericTypeNode;
use PHPStan\PhpDocParser\Ast\Type\IdentifierTypeNode;
use PHPStan\PhpDocParser\Parser\TokenIterator;
use PHPStan\Reflection\ReflectionProvider;
use PHPStan\Type\ObjectType;
use Rector\BetterPhpDocParser\PhpDocInfo\PhpDocInfoFactory;
use Rector\PhpParser\Node\BetterNodeFinder;
use Rector\Rector\AbstractRector;
use Symfony\Bundle\FrameworkBundle\Controller\AbstractController;
use Symfony\Component\Form\AbstractType;
use Symfony\Component\Form\FormFactoryInterface;
use Symfony\Component\Form\FormInterface;
use Symfony\UX\LiveComponent\Attribute\AsLiveComponent;
use Symfony\UX\LiveComponent\ComponentWithFormTrait;
use Symplify\RuleDocGenerator\ValueObject\RuleDefinition;
use Symplify\RuleDocGenerator\ValueObject\CodeSample\CodeSample;
use PHPStan\PhpDocParser\Lexer\Lexer;
use PHPStan\PhpDocParser\Parser\PhpDocParser;
use function count;
use function dd;

final class AddFormInterfaceGenericToTwigComponentWithFormRector extends AbstractRector
{
    public function __construct(
        private readonly BetterNodeFinder $betterNodeFinder,
        private readonly PhpDocInfoFactory $phpDocInfoFactory,
        private readonly ReflectionProvider $reflectionProvider,
        private readonly Lexer $lexer,
        private readonly PhpDocParser $phpDocParser,
    ) {
    }

    public function getRuleDefinition(): RuleDefinition
    {
        return new RuleDefinition(
            'Add @extends FormInterface<Foo> on instantiateForm() in Twig LiveComponent using ComponentWithFormTrait',
            [
                new CodeSample(
                    <<<'CODE'
#[AsLiveComponent]
class SomeComponent extends AbstractController
{
    use ComponentWithFormTrait;

    protected function instantiateForm(): FormInterface
    {
        return $this->createForm(FakeType::class);
    }
}
CODE,
                    <<<'CODE'
#[AsLiveComponent]
class SomeComponent extends AbstractController
{
    use ComponentWithFormTrait;

    /**
     * @extends FormInterface<FakeDto>
     */
    protected function instantiateForm(): FormInterface
    {
        return $this->createForm(FakeType::class);
    }
}
CODE
                ),
            ]
        );
    }

    public function getNodeTypes(): array
    {
        return [Class_::class];
    }

    public function refactor(Node $node): ?Node
    {
        if (! $node instanceof Class_) {
            return null;
        }

        if (! $this->hasAttribute($node, AsLiveComponent::class)) {
            return null;
        }

        if (! $this->usesTrait($node, ComponentWithFormTrait::class)) {
            return null;
        }

        $method = $node->getMethod('instantiateForm');
        if (! $method instanceof ClassMethod) {
            return null;
        }

        if ($this->phpDocInfoFactory->createFromNodeOrEmpty($method)->hasByName('return')) {
            return null;
        }

        [$formTypeClass, $dataClass] = $this->extractFormTypeClass($method);
        if (! $formTypeClass) {
            return null;
        }

        $dataClass ??= $this->resolveGenericTypeFromFormType($formTypeClass);
        if (! $dataClass) {
            return null;
        }

        $phpDocInfo = $this->phpDocInfoFactory->createFromNodeOrEmpty($method);
        $phpDocInfo->addTagValueNode(new ReturnTagValueNode(
            new GenericTypeNode(
                new IdentifierTypeNode(FormInterface::class),
                [new IdentifierTypeNode($dataClass)]
            ),
            ''
        ));

        $method->setDocComment(new Doc($phpDocInfo->getPhpDocNode()->__toString()));

        return $node;
    }

    private function hasAttribute(Class_ $class, string $attributeName): bool
    {
        foreach ($class->attrGroups as $attrGroup) {
            foreach ($attrGroup->attrs as $attr) {
                if ($this->getName($attr->name) === $attributeName) {
                    return true;
                }
            }
        }

        return false;
    }

    private function usesTrait(Class_ $class, string $traitName): bool
    {
        foreach ($class->stmts as $stmt) {
            if ($stmt instanceof TraitUse) {
                foreach ($stmt->traits as $trait) {
                    if ($this->getName($trait) === $traitName) {
                        return true;
                    }
                }
            }
        }

        return false;
    }

    private function extractFormTypeClass(ClassMethod $method): array
    {
        foreach ($method->getStmts() ?? [] as $stmt) {
            $call = $this->betterNodeFinder->findFirstInstanceOf($stmt, MethodCall::class);

            if (! $call) {
                continue;
            }

            // Method $this->createForm(SomeType::class)
            if (
                $call->var instanceof Node\Expr\Variable &&
                $this->isName($call->name, 'createForm') &&
                $this->isName($call->var, 'this')
            ) {
                $args = $call->args;
                if (isset($args[0]) && $args[0]->value instanceof ClassConstFetch) {
                    return [$this->getName($args[0]->value->class), null];
                }
            }

            // Method $this->formFactory->create(SomeType::class, $data)
            if (
                $call->var instanceof Node\Expr\PropertyFetch &&
                $call->var->var instanceof Node\Expr\Variable &&
                $this->isName($call->name, 'create') &&
                $this->isName($call->var->var, 'this')
            ) {
                $args = $call->args;

                $formTypeClass = isset($args[0]) && $args[0]->value instanceof ClassConstFetch
                    ? $this->getName($args[0]->value->class)
                    : null;

                $dataClass = null;
                if (isset($args[1])) {
                    $dataArg = $args[1]->value;
                    $dataClassType = $this->nodeTypeResolver->getType($dataArg);
                    if ($dataClassType instanceof ObjectType) {
                        $dataClass = $dataClassType->getClassName();
                    }
                }

                if ($formTypeClass !== null) {
                    return [$formTypeClass, $dataClass];
                }
            }
        }

        return [null, null];
    }

    private function resolveGenericTypeFromFormType(string $formTypeClass): ?string
    {
        if (! $this->reflectionProvider->hasClass($formTypeClass)) {
            return null;
        }

        $formTypeReflection = $this->reflectionProvider->getClass($formTypeClass);
        $reflectionClass = $formTypeReflection->getNativeReflection();

        $docComment = $reflectionClass->getDocComment();
        if (! $docComment) {
            return null;
        }

        return $this->extractGenericFromDocComment($docComment);
    }

    private function extractGenericFromDocComment(string $docComment): ?string
    {
        $tokens = new TokenIterator($this->lexer->tokenize($docComment));
        $phpDocNode = $this->phpDocParser->parse($tokens);

        foreach ($phpDocNode->children as $child) {
            if (!$child instanceof PhpDocTagNode) {
                continue;
            }

            if (!$child->value instanceof ExtendsTagValueNode) {
                continue;
            }

            $tag = $child->value;

            $type = $tag->type;

            if (!$type instanceof GenericTypeNode) {
                return null;
            }

            if ($type->type->name !== 'AbstractType') {
                return null;
            }

            if (count($type->genericTypes) !== 1) {
                return null;
            }

            $generic = $tag->type->genericTypes[0];
            if (! $generic instanceof IdentifierTypeNode) {
                return null;
            }

            return $generic->name;
        }
    }
}

tests\Rector\AddFormInterfaceGenericToTwigComponentWithFormRector\AddFormInterfaceGenericToTwigComponentWithFormRectorTest.php

<?php

declare(strict_types=1);

namespace Utils\Rector\Tests\Rector\AddFormInterfaceGenericToTwigComponentWithFormRector;

use Iterator;
use PHPUnit\Framework\Attributes\DataProvider;
use Rector\Testing\PHPUnit\AbstractRectorTestCase;

final class AddFormInterfaceGenericToTwigComponentWithFormRectorTest extends AbstractRectorTestCase
{
    #[DataProvider('provideData')]
    public function test(string $filePath): void
    {
        $this->doTestFile($filePath);
    }

    public static function provideData(): Iterator
    {
        return self::yieldFilesFromDirectory(__DIR__ . '/Fixture');
    }

    public function provideConfigFilePath(): string
    {
        return __DIR__ . '/config/configured_rule.php';
    }
}

tests\Rector\AddFormInterfaceGenericToTwigComponentWithFormRector\config\configured_rule.php

<?php

declare(strict_types=1);

use Rector\Config\RectorConfig;
use Utils\Rector\Rector\AddFormInterfaceGenericToTwigComponentWithFormRector;

return static function (RectorConfig $rectorConfig): void {
    $rectorConfig->rule(AddFormInterfaceGenericToTwigComponentWithFormRector::class);
};

tests\Rector\AddFormInterfaceGenericToTwigComponentWithFormRector\Fixture\test_fixture.php.inc

<?php

namespace Utils\Rector\Tests\Rector\AddFormInterfaceGenericToTwigComponentWithFormRector\Fixture;

use Symfony\Component\Form\FormFactoryInterface;
use Symfony\Component\Form\FormInterface;
use Symfony\UX\LiveComponent\Attribute\AsLiveComponent;
use Symfony\UX\LiveComponent\ComponentWithFormTrait;
use Symfony\UX\LiveComponent\DefaultActionTrait;
use Utils\Rector\Tests\Rector\AddFormInterfaceGenericToTwigComponentWithFormRector\Source\FakeType;
use Some\FakeDto;

#[AsLiveComponent()]
class SomeComponent
{
    use ComponentWithFormTrait;
    use DefaultActionTrait;

    private FakeDto $fakeDto;

    public function __construct(
        private readonly FormFactoryInterface $formFactory,
    ) {
    }

    protected function instantiateForm(): FormInterface
    {
        return $this->formFactory->create(FakeType::class, $this->fakeDto);
    }
}

?>
-----
<?php

namespace Utils\Rector\Tests\Rector\AddFormInterfaceGenericToTwigComponentWithFormRector\Fixture;

use Symfony\Component\Form\FormFactoryInterface;
use Symfony\Component\Form\FormInterface;
use Symfony\UX\LiveComponent\Attribute\AsLiveComponent;
use Symfony\UX\LiveComponent\ComponentWithFormTrait;
use Symfony\UX\LiveComponent\DefaultActionTrait;
use Utils\Rector\Tests\Rector\AddFormInterfaceGenericToTwigComponentWithFormRector\Source\FakeType;
use Some\FakeDto;

#[AsLiveComponent()]
class SomeComponent
{
    use ComponentWithFormTrait;
    use DefaultActionTrait;

    private FakeDto $fakeDto;

    public function __construct(
        private readonly FormFactoryInterface $formFactory,
    ) {
    }

    /**
     * @return Symfony\Component\Form\FormInterface<Some\FakeDto>
     */
    protected function instantiateForm(): FormInterface
    {
        return $this->formFactory->create(FakeType::class, $this->fakeDto);
    }
}

?>

tests\Rector\AddFormInterfaceGenericToTwigComponentWithFormRector\Fixture\test_fixture2.php.inc

<?php

namespace Utils\Rector\Tests\Rector\AddFormInterfaceGenericToTwigComponentWithFormRector\Fixture;

use Symfony\Bundle\FrameworkBundle\Controller\AbstractController;
use Symfony\Component\Form\FormInterface;
use Symfony\UX\LiveComponent\Attribute\AsLiveComponent;
use Symfony\UX\LiveComponent\ComponentWithFormTrait;
use Utils\Rector\Tests\Rector\AddFormInterfaceGenericToTwigComponentWithFormRector\Source\FakeType;

#[AsLiveComponent]
class SomeComponent extends AbstractController
{
    use ComponentWithFormTrait;

    protected function instantiateForm(): FormInterface
    {
        return $this->createForm(FakeType::class);
    }
}

?>
-----
<?php

namespace Utils\Rector\Tests\Rector\AddFormInterfaceGenericToTwigComponentWithFormRector\Fixture;

use Symfony\Bundle\FrameworkBundle\Controller\AbstractController;
use Symfony\Component\Form\FormInterface;
use Symfony\UX\LiveComponent\Attribute\AsLiveComponent;
use Symfony\UX\LiveComponent\ComponentWithFormTrait;
use Utils\Rector\Tests\Rector\AddFormInterfaceGenericToTwigComponentWithFormRector\Source\FakeType;

#[AsLiveComponent]
class SomeComponent extends AbstractController
{
    use ComponentWithFormTrait;

    /**
     * @return Symfony\Component\Form\FormInterface<Some\FakeDto>
     */
    protected function instantiateForm(): FormInterface
    {
        return $this->createForm(FakeType::class);
    }
}

?>

tests\Rector\AddFormInterfaceGenericToTwigComponentWithFormRector\Source\FakeType.php

<?php

declare(strict_types=1);

namespace Utils\Rector\Tests\Rector\AddFormInterfaceGenericToTwigComponentWithFormRector\Source;

use Symfony\Component\Form\AbstractType;

/**
 * @extends AbstractType<Some\FakeDto>
 */
final class FakeType extends AbstractType
{

}