added TraitUse, represents use-statement

Author: dg

readme.md

@@ -53,7 +53,6 @@ $class
 	->setFinal()
 	->setExtends(ParentClass::class)
 	->addImplement(Countable::class)
-	->addTrait(Nette\SmartObject::class)
 	->addComment("Description of class.\nSecond line\n")
 	->addComment('@property-read Nette\Forms\Form $form');
 
@@ -72,7 +71,6 @@ It will render this result:
  */
 final class Demo extends ParentClass implements Countable
 {
-	use Nette\SmartObject;
 }
 ```
 
@@ -315,7 +313,8 @@ Using Traits
 ```php
 $class = new Nette\PhpGenerator\ClassType('Demo');
 $class->addTrait('SmartObject');
-$class->addTrait('MyTrait', ['sayHello as protected']);
+$class->addTrait('MyTrait', true)
+	->addResolution('sayHello as protected');
 echo $class;
 ```
 

src/PhpGenerator/ClassType.php

@@ -56,7 +56,7 @@ final class ClassType
 	/** @var string[] */
 	private $implements = [];
 
-	/** @var array[] */
+	/** @var TraitUse[] */
 	private $traits = [];
 
 	/** @var Constant[] name => Constant */
@@ -330,13 +330,18 @@ public function removeImplement(string $name): self
 
 
 	/**
-	 * @param  string[]  $names
+	 * @param  string[]|TraitUse[]  $traits
 	 * @return static
 	 */
-	public function setTraits(array $names): self
+	public function setTraits(array $traits): self
 	{
-		$this->validateNames($names);
-		$this->traits = array_fill_keys($names, []);
+		$this->traits = [];
+		foreach ($traits as $trait) {
+			if (!$trait instanceof TraitUse) {
+				$trait = new TraitUse($trait);
+			}
+			$this->traits[$trait->getName()] = $trait;
+		}
 		return $this;
 	}
 
@@ -355,11 +360,19 @@ public function getTraitResolutions(): array
 	}
 
 
-	/** @return static */
-	public function addTrait(string $name, array $resolutions = []): self
+	/**
+	 * @param  array|bool  $resolutions
+	 * @return static|TraitUse
+	 */
+	public function addTrait(string $name, $resolutions = [])
 	{
-		$this->validateNames([$name]);
-		$this->traits[$name] = $resolutions;
+		$this->traits[$name] = $trait = new TraitUse($name);
+		if ($resolutions === true) {
+			return $trait;
+		}
+		array_map(function ($item) use ($trait) {
+			$trait->addResolution($item);
+		}, $resolutions);
 		return $this;
 	}
 
@@ -373,7 +386,7 @@ public function removeTrait(string $name): self
 
 
 	/**
-	 * @param  Method|Property|Constant|EnumCase  $member
+	 * @param  Method|Property|Constant|EnumCase|TraitUse  $member
 	 * @return static
 	 */
 	public function addMember($member): self
@@ -393,8 +406,11 @@ public function addMember($member): self
 		} elseif ($member instanceof EnumCase) {
 			$this->cases[$member->getName()] = $member;
 
+		} elseif ($member instanceof TraitUse) {
+			$this->traits[$member->getName()] = $member;
+
 		} else {
-			throw new Nette\InvalidArgumentException('Argument must be Method|Property|Constant|EnumCase.');
+			throw new Nette\InvalidArgumentException('Argument must be Method|Property|Constant|EnumCase|TraitUse.');
 		}
 
 		return $this;

src/PhpGenerator/Printer.php

@@ -140,8 +140,9 @@ public function printClass(ClassType $class, PhpNamespace $namespace = null): st
 			: function ($s) { return $s; };
 
 		$traits = [];
-		foreach ($class->getTraitResolutions() as $trait => $resolutions) {
-			$traits[] = 'use ' . $resolver($trait)
+		foreach ($class->getTraitResolutions() as $trait) {
+			$resolutions = $trait->getResolutions();
+			$traits[] = 'use ' . $resolver($trait->getName())
 				. ($resolutions ? " {\n" . $this->indentation . implode(";\n" . $this->indentation, $resolutions) . ";\n}\n" : ";\n");
 		}
 

src/PhpGenerator/TraitUse.php

@@ -0,0 +1,47 @@
+<?php
+
+/**
+ * This file is part of the Nette Framework (https://nette.org)
+ * Copyright (c) 2004 David Grudl (https://davidgrudl.com)
+ */
+
+declare(strict_types=1);
+
+namespace Nette\PhpGenerator;
+
+use Nette;
+
+
+/**
+ * use Trait
+ */
+final class TraitUse
+{
+	use Nette\SmartObject;
+	use Traits\NameAware;
+
+	/** @var array */
+	private $resolutions = [];
+
+
+	public function __construct(string $name)
+	{
+		if (!Nette\PhpGenerator\Helpers::isNamespaceIdentifier($name, true)) {
+			throw new Nette\InvalidArgumentException("Value '$name' is not valid trait name.");
+		}
+		$this->name = $name;
+	}
+
+
+	public function addResolution(string $resolution): self
+	{
+		$this->resolutions[] = $resolution;
+		return $this;
+	}
+
+
+	public function getResolutions(): array
+	{
+		return $this->resolutions;
+	}
+}

tests/PhpGenerator/ClassType.addMember.phpt

@@ -12,17 +12,19 @@ require __DIR__ . '/../bootstrap.php';
 Assert::exception(function () {
 	(new ClassType('Example'))
 		->addMember(new stdClass);
-}, Nette\InvalidArgumentException::class, 'Argument must be Method|Property|Constant|EnumCase.');
+}, Nette\InvalidArgumentException::class, 'Argument must be Method|Property|Constant|EnumCase|TraitUse.');
 
 
 $class = (new ClassType('Example'))
 	->addMember($method = new Nette\PhpGenerator\Method('getHandle'))
 	->addMember($property = new Nette\PhpGenerator\Property('handle'))
-	->addMember($const = new Nette\PhpGenerator\Constant('ROLE'));
+	->addMember($const = new Nette\PhpGenerator\Constant('ROLE'))
+	->addMember($trait = new Nette\PhpGenerator\TraitUse('Foo\Bar'));
 
 Assert::same(['getHandle' => $method], $class->getMethods());
 Assert::same(['handle' => $property], $class->getProperties());
 Assert::same(['ROLE' => $const], $class->getConstants());
+Assert::same(['Foo\Bar'], $class->getTraits());
 Assert::same('', $method->getBody());
 
 

tests/PhpGenerator/ClassType.phpt

@@ -43,7 +43,6 @@ Assert::false($class->isFinal());
 Assert::true($class->isAbstract());
 Assert::same('ParentClass', $class->getExtends());
 Assert::same(['ObjectTrait', 'AnotherTrait'], $class->getTraits());
-Assert::same(['ObjectTrait' => [], 'AnotherTrait' => ['sayHello as protected']], $class->getTraitResolutions());
 Assert::count(2, $class->getConstants());
 Assert::type(Nette\PhpGenerator\Constant::class, $class->getConstants()['ROLE']);
 
@@ -139,6 +138,11 @@ $class->removeTrait('foo');
 $class->addImplement('foo');
 $class->removeImplement('foo');
 
+$class
+	->addTrait('ThirdTrait', true)
+		->addResolution('a as private foo')
+		->addResolution('b as private bar');
+
 
 sameFile(__DIR__ . '/expected/ClassType.expect', (string) $class);
 

tests/PhpGenerator/expected/ClassType.expect

@@ -10,6 +10,10 @@ abstract class Example extends ParentClass implements IExample, IOne
 	use AnotherTrait {
 		sayHello as protected;
 	}
+	use ThirdTrait {
+		a as private foo;
+		b as private bar;
+	}
 
 	const ROLE = 'admin';
 	final public const ACTIVE = false;

tests/PhpGenerator/invalidNames.phpt

@@ -110,11 +110,11 @@ Assert::exception(function () use ($class) {
 
 Assert::exception(function () use ($class) {
 	$class->setTraits(['A', '*']);
-}, Nette\InvalidArgumentException::class, "Value '*' is not valid class name.");
+}, Nette\InvalidArgumentException::class, "Value '*' is not valid trait name.");
 
 Assert::exception(function () use ($class) {
 	$class->addTrait('*');
-}, Nette\InvalidArgumentException::class, "Value '*' is not valid class name.");
+}, Nette\InvalidArgumentException::class, "Value '*' is not valid trait name.");
 
 
 Assert::noError(function () {