Add trait support for class builder - Close #13

Author: sandrokeil

src/Builder/ClassBuilder.php

@@ -15,6 +15,7 @@
 use OpenCodeModeling\CodeAst\NodeVisitor\ClassFile;
 use OpenCodeModeling\CodeAst\NodeVisitor\ClassImplements;
 use OpenCodeModeling\CodeAst\NodeVisitor\ClassNamespace;
+use OpenCodeModeling\CodeAst\NodeVisitor\ClassUseTrait;
 use OpenCodeModeling\CodeAst\NodeVisitor\NamespaceUse;
 use OpenCodeModeling\CodeAst\NodeVisitor\StrictType;
 use PhpParser\Node;
@@ -50,6 +51,9 @@ final class ClassBuilder
     /** @var string[] */
     private $namespaceUse = [];
 
+    /** @var string[] */
+    private $traits = [];
+
     /** @var ClassConstBuilder[] */
     private $constants = [];
 
@@ -125,6 +129,13 @@ public function setNamespaceUse(string ...$namespaces): self
         return $this;
     }
 
+    public function setUseTrait(string ...$traits): self
+    {
+        $this->traits = $traits;
+
+        return $this;
+    }
+
     public function setConstants(ClassConstBuilder ...$constants): self
     {
         $this->constants = $constants;
@@ -183,6 +194,14 @@ public function getNamespaceUse(): array
         return $this->namespaceUse;
     }
 
+    /**
+     * @return string[]
+     */
+    public function getUseTrait(): array
+    {
+        return $this->traits;
+    }
+
     /**
      * @return ClassConstBuilder[]
      */
@@ -218,6 +237,10 @@ public function generate(): array
         if ($this->implements) {
             $visitors[] = new ClassImplements(...$this->implements);
         }
+        if ($this->traits) {
+            $visitors[] = new ClassUseTrait(...$this->traits);
+        }
+
         if (\count($this->constants) > 0) {
             \array_push(
                 $visitors,
@@ -254,12 +277,19 @@ private function unpackNode(Node $node): void
                 }
                 break;
             case $node instanceof Node\Stmt\UseUse:
-                $this->namespaceUse[] = $node->name->toString();
+                $this->namespaceUse[] = $node->name instanceof Node\Name\FullyQualified
+                    ? '\\' . $node->name->toString()
+                    : $node->name->toString();
                 break;
             case $node instanceof Node\Stmt\Class_:
                 $this->name = $node->name->name;
                 $this->final = $node->isFinal();
-                $this->extends = $node->extends ? $node->extends->toString() : null;
+
+                if ($node->extends !== null) {
+                    $this->extends = $node->extends instanceof Node\Name\FullyQualified
+                        ? '\\' . $node->extends->toString()
+                        : $node->extends->toString();
+                }
 
                 foreach ($node->stmts as $stmt) {
                     $this->unpackNode($stmt);
@@ -273,6 +303,16 @@ static function (Node\Name $name) {
                     $node->implements
                 );
                 break;
+            case $node instanceof Node\Stmt\TraitUse:
+                $this->traits = \array_map(
+                    static function (Node\Name $name) {
+                        return $name instanceof Node\Name\FullyQualified
+                            ? '\\' . $name->toString()
+                            : $name->toString();
+                    },
+                    $node->traits
+                );
+                break;
             case $node instanceof Node\Stmt\ClassConst:
                 $this->constants[] = ClassConstBuilder::fromNode($node);
                 break;

tests/Builder/ClassBuilderTest.php

@@ -43,6 +43,7 @@ public function it_generates_class_for_empty_file(): void
             ->setExtends('BaseClass')
             ->setNamespaceUse('Foo\\Bar')
             ->setImplements('\\Iterator', 'Bar')
+            ->setUseTrait('\\My\\TestTrait')
             ->setConstants(ClassConstBuilder::fromScratch('PRIV', 'private')->setPrivate());
 
         $nodeTraverser = new NodeTraverser();
@@ -57,6 +58,7 @@ public function it_generates_class_for_empty_file(): void
 use Foo\Bar;
 final class TestClass extends BaseClass implements \Iterator, Bar
 {
+    use \My\TestTrait;
     private const PRIV = 'private';
 }
 EOF;
@@ -78,6 +80,7 @@ public function it_generates_class_for_empty_file_from_template(): void
 use Foo\Bar;
 final class TestClass extends BaseClass implements \Iterator, Bar
 {
+    use \My\TestTrait;
     const FIRST = 1;
     private const PRIV = 'private';
     protected const PROT = 'protected';
@@ -94,6 +97,7 @@ final class TestClass extends BaseClass implements \Iterator, Bar
         $this->assertTrue($classFactory->isFinal());
         $this->assertTrue($classFactory->isStrict());
         $this->assertFalse($classFactory->isAbstract());
+        $this->assertSame(['\\My\\TestTrait'], $classFactory->getUseTrait());
 
         $nodeTraverser = new NodeTraverser();
         $classFactory->injectVisitors($nodeTraverser);