Add support for adding custom traits to models

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);

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)))));

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;

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'],
         ];
     }
 

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);

tests/fixtures/drafts/model-with-traits.yaml

@@ -0,0 +1,4 @@
+models:
+  User:
+    traits:
+      - Spatie\Permission\Traits\HasRoles

tests/fixtures/models/model-with-traits.php

@@ -0,0 +1,28 @@
+<?php
+
+namespace App\Models;
+
+use Illuminate\Database\Eloquent\Factories\HasFactory;
+use Illuminate\Database\Eloquent\Model;
+use Spatie\Permission\Traits\HasRoles;
+
+class User extends Model
+{
+    use HasFactory, HasRoles;
+
+    /**
+     * The attributes that are mass assignable.
+     *
+     * @var array
+     */
+    protected $fillable = [];
+
+    /**
+     * The attributes that should be cast to native types.
+     *
+     * @var array
+     */
+    protected $casts = [
+        'id' => 'integer',
+    ];
+}