diff --git a/src/Generators/ModelGenerator.php b/src/Generators/ModelGenerator.php index 6e7f578f..d9d98d4b 100644 --- a/src/Generators/ModelGenerator.php +++ b/src/Generators/ModelGenerator.php @@ -460,6 +460,13 @@ protected function addTraits(Model $model, $stub): string $traits[] = 'HasUuids'; } + if ($model->usesCustomTraits()) { + foreach ($model->customTraits() as $trait) { + $this->addImport($model, $trait); + $traits[] = Str::afterLast($trait, '\\'); + } + } + sort($traits); return Str::replaceFirst('use HasFactory', 'use ' . implode(', ', $traits), $stub); diff --git a/src/Lexers/ModelLexer.php b/src/Lexers/ModelLexer.php index 8ec44a64..e4b46dec 100644 --- a/src/Lexers/ModelLexer.php +++ b/src/Lexers/ModelLexer.php @@ -194,6 +194,14 @@ private function buildModel(string $name, array $columns): Model unset($columns['relationships']); } + if (isset($columns['traits'])) { + foreach (explode(' ', $columns['traits']) as $trait) { + $model->addCustomTrait(trim($trait)); + } + + unset($columns['traits']); + } + if (isset($columns['indexes'])) { foreach ($columns['indexes'] as $index) { $model->addIndex(new Index(key($index), array_map('trim', explode(',', current($index))))); diff --git a/src/Models/Model.php b/src/Models/Model.php index 207e2458..26f8b374 100644 --- a/src/Models/Model.php +++ b/src/Models/Model.php @@ -33,6 +33,8 @@ class Model implements BlueprintModel private array $indexes = []; + private array $customTraits = []; + public function __construct($name) { $this->name = class_basename($name); @@ -111,6 +113,11 @@ public function usesUuids(): bool return $this->usesPrimaryKey() && $this->columns[$this->primaryKey]->dataType() === 'uuid'; } + public function usesCustomTraits(): bool + { + return count($this->customTraits) > 0; + } + public function idType(): ?string { if (!$this->usesPrimaryKey()) { @@ -252,6 +259,16 @@ public function addIndex(Index $index): void $this->indexes[] = $index; } + public function customTraits(): array + { + return $this->customTraits; + } + + public function addCustomTrait(string $trait): void + { + $this->customTraits[] = $trait; + } + public function pivotTables(): array { return $this->pivotTables; diff --git a/tests/Feature/Generators/ModelGeneratorTest.php b/tests/Feature/Generators/ModelGeneratorTest.php index 40870e39..3d4cc5d3 100644 --- a/tests/Feature/Generators/ModelGeneratorTest.php +++ b/tests/Feature/Generators/ModelGeneratorTest.php @@ -80,6 +80,7 @@ public function output_generates_models($definition, $path, $model): void $this->filesystem->expects('exists') ->with(dirname($path)) ->andReturnTrue(); + $this->filesystem->expects('put') ->with($path, $this->fixture($model)); @@ -648,6 +649,7 @@ public static function modelTreeDataProvider(): array ['drafts/infer-belongsto.yaml', 'app/Models/Conference.php', 'models/infer-belongsto.php'], ['drafts/model-with-ulid-id.yaml', 'app/Models/User.php', 'models/model-with-ulid-trait.php'], ['drafts/model-with-uuid-id.yaml', 'app/Models/User.php', 'models/model-with-uuid-trait.php'], + ['drafts/model-with-traits.yaml', 'app/Models/User.php', 'models/model-with-traits.php'], ]; } diff --git a/tests/Feature/Lexers/ModelLexerTest.php b/tests/Feature/Lexers/ModelLexerTest.php index 1d1abb07..af263c54 100644 --- a/tests/Feature/Lexers/ModelLexerTest.php +++ b/tests/Feature/Lexers/ModelLexerTest.php @@ -42,6 +42,7 @@ public function it_returns_models(): void ], 'ModelThree' => [ 'id' => 'increments', + 'traits' => 'Spatie\Permission\Traits\HasRoles', ], ], ]; @@ -55,6 +56,7 @@ public function it_returns_models(): void $this->assertEquals('ModelOne', $model->name()); $this->assertTrue($model->usesTimestamps()); $this->assertFalse($model->usesSoftDeletes()); + $this->assertFalse($model->usesCustomTraits()); $columns = $model->columns(); $this->assertCount(2, $columns); @@ -69,6 +71,8 @@ public function it_returns_models(): void $this->assertEquals('ModelTwo', $model->name()); $this->assertTrue($model->usesTimestamps()); $this->assertFalse($model->usesSoftDeletes()); + $this->assertFalse($model->usesCustomTraits()); + $columns = $model->columns(); $this->assertCount(2, $columns); @@ -83,6 +87,7 @@ public function it_returns_models(): void $this->assertEquals('ModelThree', $model->name()); $this->assertTrue($model->usesTimestamps()); $this->assertFalse($model->usesSoftDeletes()); + $this->assertTrue($model->usesCustomTraits()); $columns = $model->columns(); $this->assertCount(1, $columns); diff --git a/tests/fixtures/drafts/model-with-traits.yaml b/tests/fixtures/drafts/model-with-traits.yaml new file mode 100644 index 00000000..9bf191db --- /dev/null +++ b/tests/fixtures/drafts/model-with-traits.yaml @@ -0,0 +1,4 @@ +models: + User: + traits: + - Spatie\Permission\Traits\HasRoles diff --git a/tests/fixtures/models/model-with-traits.php b/tests/fixtures/models/model-with-traits.php new file mode 100644 index 00000000..d117b155 --- /dev/null +++ b/tests/fixtures/models/model-with-traits.php @@ -0,0 +1,28 @@ + 'integer', + ]; +}