Support custom node visitors in class and interface builder - Close #65

Author: sandrokeil

src/Builder/ClassBuilder.php

@@ -64,6 +64,9 @@ final class ClassBuilder implements File
     /** @var ClassMethodBuilder[] */
     private array $methods = [];
 
+    /** @var NodeVisitor[] */
+    private array $nodeVisitors = [];
+
     private function __construct()
     {
     }
@@ -414,6 +417,18 @@ public function hasMethod(string $methodName): bool
         return isset($this->methods[$methodName]);
     }
 
+    public function setNodeVisitors(NodeVisitor ...$nodeVisitors): void
+    {
+        $this->nodeVisitors = $nodeVisitors;
+    }
+
+    public function addNodeVisitor(NodeVisitor ...$nodeVisitors): void
+    {
+        foreach ($nodeVisitors as $nodeVisitor) {
+            $this->nodeVisitors[] = $nodeVisitor;
+        }
+    }
+
     public function getNamespace(): ?string
     {
         return $this->namespace;
@@ -649,7 +664,7 @@ static function (ClassMethodBuilder $method) use ($parser) {
             );
         }
 
-        return $visitors;
+        return \array_merge($visitors, $this->nodeVisitors);
     }
 
     private function unpackNode(Node $node): void

src/Builder/InterfaceBuilder.php

@@ -47,6 +47,9 @@ final class InterfaceBuilder implements File
     /** @var ClassMethodBuilder[] */
     private array $methods = [];
 
+    /** @var NodeVisitor[] */
+    private array $nodeVisitors = [];
+
     private function __construct()
     {
     }
@@ -265,6 +268,18 @@ public function hasMethod(string $methodName): bool
         return isset($this->methods[$methodName]);
     }
 
+    public function setNodeVisitors(NodeVisitor ...$nodeVisitors): void
+    {
+        $this->nodeVisitors = $nodeVisitors;
+    }
+
+    public function addNodeVisitor(NodeVisitor ...$nodeVisitors): void
+    {
+        foreach ($nodeVisitors as $nodeVisitor) {
+            $this->nodeVisitors[] = $nodeVisitor;
+        }
+    }
+
     /**
      * Uses uasort internally
      *
@@ -402,7 +417,7 @@ static function (ClassMethodBuilder $method) use ($parser) {
             );
         }
 
-        return $visitors;
+        return \array_merge($visitors, $this->nodeVisitors);
     }
 
     private function unpackNode(Node $node): void

tests/Builder/ClassBuilderTest.php

@@ -12,7 +12,11 @@
 
 use OpenCodeModeling\CodeAst\Builder\ClassBuilder;
 use OpenCodeModeling\CodeAst\Builder\ClassConstBuilder;
+use OpenCodeModeling\CodeAst\Builder\ClassMethodBuilder;
+use PhpParser\Node;
 use PhpParser\NodeTraverser;
+use PhpParser\NodeVisitor;
+use PhpParser\NodeVisitorAbstract;
 use PhpParser\Parser;
 use PhpParser\ParserFactory;
 use PhpParser\PrettyPrinter\Standard;
@@ -462,4 +466,75 @@ final class TestClass
 
         $this->assertSame($expected, $this->printer->prettyPrintFile($nodeTraverser->traverse($this->parser->parse(''))));
     }
+
+    /**
+     * @test
+     */
+    public function it_supports_adding_of_node_visitors(): void
+    {
+        $ast = $this->parser->parse('');
+
+        $classFactory = ClassBuilder::fromScratch('TestClass', 'My\\Awesome\\Service');
+        $classFactory->setMethods(
+            ClassMethodBuilder::fromScratch('setActive')->setReturnType('void')->setStatic(true)
+        );
+
+        $classFactory->addNodeVisitor($this->getSetActiveNodeVisitor());
+
+        $nodeTraverser = new NodeTraverser();
+        $classFactory->injectVisitors($nodeTraverser, $this->parser);
+
+        $expected = <<<'EOF'
+<?php
+
+declare (strict_types=1);
+namespace My\Awesome\Service;
+
+class TestClass
+{
+    public static function setActive() : void
+    {
+        $tmp = $this->get();
+    }
+}
+EOF;
+
+        $this->assertSame($expected, $this->printer->prettyPrintFile($nodeTraverser->traverse($ast)));
+    }
+
+    private function getSetActiveNodeVisitor(): NodeVisitor
+    {
+        $nodes = $this->parser->parse('<?php $tmp = $this->get();');
+
+        return new class($nodes) extends NodeVisitorAbstract {
+            private $nodes;
+
+            public function __construct($nodes)
+            {
+                $this->nodes = $nodes;
+            }
+
+            public function afterTraverse(array $nodes)
+            {
+                $newNodes = [];
+
+                foreach ($nodes as $node) {
+                    $newNodes[] = $node;
+
+                    if (! $node instanceof Node\Stmt\Class_) {
+                        continue;
+                    }
+
+                    if ($node->stmts[0] instanceof Node\Stmt\ClassMethod
+                        && $node->stmts[0]->name instanceof Node\Identifier
+                        && $node->stmts[0]->name->name === 'setActive'
+                    ) {
+                        $node->stmts[0]->stmts = $this->nodes;
+                    }
+                }
+
+                return $newNodes;
+            }
+        };
+    }
 }