Allow TraitUse statements in trait builder

Fixes #413.

CHANGELOG.md

@@ -4,6 +4,7 @@ Version 3.1.1-dev
 ### Fixed
 
 * Fixed syntax error on comment after brace-style namespace declaration. (#412)
+* Added support for TraitUse statements in trait builder. (#413)
 
 Version 3.1.0 (2017-07-28)
 --------------------------

lib/PhpParser/Builder/Trait_.php

@@ -8,6 +8,7 @@ use PhpParser\Node\Stmt;
 class Trait_ extends Declaration
 {
     protected $name;
+    protected $uses = array();
     protected $properties = array();
     protected $methods = array();
 
@@ -34,6 +35,8 @@ class Trait_ extends Declaration
             $this->properties[] = $stmt;
         } else if ($stmt instanceof Stmt\ClassMethod) {
             $this->methods[] = $stmt;
+        } else if ($stmt instanceof Stmt\TraitUse) {
+            $this->uses[] = $stmt;
         } else {
             throw new \LogicException(sprintf('Unexpected node of type "%s"', $stmt->getType()));
         }
@@ -49,7 +52,7 @@ class Trait_ extends Declaration
     public function getNode() {
         return new Stmt\Trait_(
             $this->name, array(
-                'stmts' => array_merge($this->properties, $this->methods)
+                'stmts' => array_merge($this->uses, $this->properties, $this->methods)
             ), $this->attributes
         );
     }

test/PhpParser/Builder/TraitTest.php

@@ -3,6 +3,7 @@
 namespace PhpParser\Builder;
 
 use PhpParser\Comment;
+use PhpParser\Node\Name;
 use PhpParser\Node\Stmt;
 
 class TraitTest extends \PHPUnit_Framework_TestCase
@@ -18,19 +19,21 @@ class TraitTest extends \PHPUnit_Framework_TestCase
         $prop = new Stmt\Property(Stmt\Class_::MODIFIER_PUBLIC, array(
             new Stmt\PropertyProperty('test')
         ));
+        $use = new Stmt\TraitUse([new Name('OtherTrait')]);
         $trait = $this->createTraitBuilder('TestTrait')
             ->setDocComment('/** Nice trait */')
             ->addStmt($method1)
-            ->addStmts(array($method2, $method3))
+            ->addStmts([$method2, $method3])
             ->addStmt($prop)
+            ->addStmt($use)
             ->getNode();
-        $this->assertEquals(new Stmt\Trait_('TestTrait', array(
-            'stmts' => array($prop, $method1, $method2, $method3)
-        ), array(
-            'comments' => array(
+        $this->assertEquals(new Stmt\Trait_('TestTrait', [
+            'stmts' => [$use, $prop, $method1, $method2, $method3]
+        ], [
+            'comments' => [
                 new Comment\Doc('/** Nice trait */')
-            )
-        )), $trait);
+            ]
+        ]), $trait);
     }
 
     /**