Add trait handling

app/Tools/File.php

@@ -5,6 +5,7 @@
 use Exception;
 use Illuminate\Support\Facades\Storage;
 use Tighten\Mise\Tools\PhpParser\Visitors\AddImportVisitor;
+use Tighten\Mise\Tools\PhpParser\Visitors\AddTraitVisitor;
 
 class File extends ConsoleCommand
 {
@@ -191,6 +192,14 @@ public function addImports(string $path, string|array $classes): static
         return $this;
     }
 
+    public function addTraits(string $path, string|array $traits): static
+    {
+        (new PhpParser)->edit($path, [new AddTraitVisitor($traits)]);
+        (new CsFixer)->fix($path, ['class_attributes_separation']);
+
+        return $this;
+    }
+
     /**
      * Allow passing globbing patterns or arrays of globbing patterns.
      */

app/Tools/PhpParser.php

@@ -6,6 +6,7 @@
 use Illuminate\Support\Facades\Cache;
 use Illuminate\Support\Facades\Storage;
 use PhpParser\NodeTraverser;
+use PhpParser\NodeVisitor\CloningVisitor;
 use PhpParser\NodeVisitorAbstract;
 use PhpParser\Parser;
 use PhpParser\ParserFactory;
@@ -28,6 +29,7 @@ public function __construct()
 
     /**
      * Central entry point for editing PHP files with PHP Parser.
+     * Uses formatting-preserving pretty printing to maintain original code style.
      *
      * @param  array<NodeVisitorAbstract>  $edits
      *
@@ -37,10 +39,24 @@ public function edit(string $phpFilePath, array $edits): void
     {
         Cache::put(self::ACTIVE_FILENAME_KEY, $phpFilePath);
 
-        $this->traverser($edits)
-            ->traverse($ast = $this->toAst($phpFilePath));
+        $code = Storage::get($phpFilePath);
 
-        Storage::put($phpFilePath, $this->toPhpFile($ast));
+        // Parse the code and save the tokens for format preservation
+        $oldStmts = $this->parser->parse($code);
+        $oldTokens = $this->parser->getTokens();
+
+        // Clone the AST before making changes
+        $cloningTraverser = new NodeTraverser;
+        $cloningTraverser->addVisitor(new CloningVisitor);
+        $newStmts = $cloningTraverser->traverse($oldStmts);
+
+        // Apply the edits to the cloned AST
+        $this->traverser($edits)->traverse($newStmts);
+
+        // Use format-preserving pretty printing
+        $newCode = $this->phpFilePrinter->printFormatPreserving($newStmts, $oldStmts, $oldTokens);
+
+        Storage::put($phpFilePath, $newCode);
     }
 
     /**

app/Tools/PhpParser/Visitors/AddTraitVisitor.php

@@ -0,0 +1,200 @@
+<?php
+
+namespace Tighten\Mise\Tools\PhpParser\Visitors;
+
+use Exception;
+use Illuminate\Support\Facades\Cache;
+use PhpParser\Node;
+use PhpParser\Node\Name;
+use PhpParser\Node\Stmt\Class_;
+use PhpParser\Node\Stmt\TraitUse;
+use PhpParser\NodeVisitorAbstract;
+use Tighten\Mise\Tools\PhpParser;
+
+class AddTraitVisitor extends NodeVisitorAbstract
+{
+    use InteractsWithNodes;
+
+    /** @var array<string> */
+    private array $traits;
+
+    /** @param string|array<string> $traits */
+    public function __construct(string|array $traits)
+    {
+        $this->traits = is_array($traits) ? $traits : [$traits];
+    }
+
+    /**
+     * Validate the document before traversing it.
+     *
+     * @param  array<Node>  $nodes
+     *
+     * @throws Exception
+     */
+    public function beforeTraverse(array $nodes): ?array
+    {
+        if (! $this->hasClass($nodes)) {
+            throw new Exception('Class not found in ' . Cache::get(PhpParser::ACTIVE_FILENAME_KEY));
+        }
+
+        return $nodes;
+    }
+
+    /**
+     * Parse a node after traversing it.
+     * If it's a class, add the traits.
+     */
+    public function enterNode(Node $node)
+    {
+        if ($node instanceof Class_) {
+            $this->addTraitsToClass($node);
+        }
+    }
+
+    /**
+     * Add traits to a class node.
+     */
+    private function addTraitsToClass(Class_ $class): void
+    {
+        $existingTraitUseInfo = $this->findExistingTraitUse($class);
+        $existingTraits = $this->extractExistingTraitNames($existingTraitUseInfo);
+        $traitsToAdd = $this->determineTraitsToAdd($existingTraits);
+
+        if (empty($traitsToAdd)) {
+            return;
+        }
+
+        if ($existingTraitUseInfo['statement'] !== null) {
+            $this->updateExistingTraitUse($class, $existingTraitUseInfo, $traitsToAdd);
+        } else {
+            $this->createNewTraitUse($class, $traitsToAdd);
+        }
+    }
+
+    /**
+     * Find the existing trait use statement in a class.
+     *
+     * @return array{statement: ?TraitUse, index: ?int}
+     */
+    private function findExistingTraitUse(Class_ $class): array
+    {
+        foreach ($class->stmts as $index => $stmt) {
+            if ($stmt instanceof TraitUse) {
+                return [
+                    'statement' => $stmt,
+                    'index' => $index,
+                ];
+            }
+        }
+
+        return [
+            'statement' => null,
+            'index' => null,
+        ];
+    }
+
+    /**
+     * Extract the simple names of existing traits.
+     *
+     * @param  array{statement: ?TraitUse, index: ?int}  $traitUseInfo
+     * @return array<string>
+     */
+    private function extractExistingTraitNames(array $traitUseInfo): array
+    {
+        if ($traitUseInfo['statement'] === null) {
+            return [];
+        }
+
+        $existingTraits = [];
+        foreach ($traitUseInfo['statement']->traits as $trait) {
+            $existingTraits[] = $this->getSimpleName($this->getTraitName($trait));
+        }
+
+        return $existingTraits;
+    }
+
+    /**
+     * Determine which traits need to be added.
+     *
+     * @param  array<string>  $existingTraits
+     * @return array<string>
+     */
+    private function determineTraitsToAdd(array $existingTraits): array
+    {
+        $traitsToAdd = [];
+
+        foreach ($this->traits as $trait) {
+            $traitSimpleName = $this->getSimpleName($trait);
+            if (! in_array($traitSimpleName, $existingTraits)) {
+                $traitsToAdd[] = $trait;
+            }
+        }
+
+        return $traitsToAdd;
+    }
+
+    /**
+     * Update an existing trait use statement with new traits.
+     *
+     * @param  array{statement: TraitUse, index: int}  $existingTraitUseInfo
+     * @param  array<string>  $traitsToAdd
+     */
+    private function updateExistingTraitUse(Class_ $class, array $existingTraitUseInfo, array $traitsToAdd): void
+    {
+        // Combine existing traits with new ones
+        $allTraits = $existingTraitUseInfo['statement']->traits;
+        foreach ($traitsToAdd as $trait) {
+            $allTraits[] = new Name($trait);
+        }
+
+        // Create new TraitUse statement, preserving attributes (including comments) from the original
+        $newTraitUse = new TraitUse($allTraits);
+        $newTraitUse->setAttributes($existingTraitUseInfo['statement']->getAttributes());
+
+        // Replace the existing trait use statement
+        $class->stmts[$existingTraitUseInfo['index']] = $newTraitUse;
+    }
+
+    /**
+     * Create a new trait use statement and add it to the class.
+     *
+     * @param  array<string>  $traitsToAdd
+     */
+    private function createNewTraitUse(Class_ $class, array $traitsToAdd): void
+    {
+        $traitNodes = array_map(fn ($trait) => new Name($trait), $traitsToAdd);
+        $newTraitUse = new TraitUse($traitNodes);
+
+        // Insert at the beginning of the class body
+        array_unshift($class->stmts, $newTraitUse);
+    }
+
+    /**
+     * Check if the AST has a class declaration.
+     *
+     * @param  array<Node>  $ast
+     */
+    private function hasClass(array $ast): bool
+    {
+        return collect($this->getStatements($ast))
+            ->contains(fn ($value) => $value instanceof Class_);
+    }
+
+    /**
+     * Get the trait name from a Name node.
+     */
+    private function getTraitName(Name $name): string
+    {
+        return $name->toString();
+    }
+
+    /**
+     * Get the simple name (without namespace) from a fully qualified trait name.
+     */
+    private function getSimpleName(string $trait): string
+    {
+        $parts = explode('\\', $trait);
+
+        return end($parts);
+    }
+}

tests/Feature/Tools/FileTest.php

@@ -218,7 +218,7 @@
 
     (new File)->addImports($path, 'App\Models\User');
 
-    expect(Storage::get($path))->toBe("<?php\n\nnamespace App\Awesome;\n\nuse App\Models\Contact;\nuse App\Models\User;\n\nclass Test\n{\n    // Some code\n}");
+    expect(Storage::get($path))->toBe("<?php\n\nnamespace App\Awesome;\n\nuse App\Models\Contact;\nuse App\Models\User;\n\nclass Test {\n    // Some code\n}");
 });
 
 test('file->addImport(...) skips duplicate imports', function () {
@@ -294,3 +294,94 @@
 
     expect(Storage::get($path))->toBe("Line 1\n    Line 2\n    Appended content\n");
 });
+
+test('file->addTraits(...) adds trait to class with no traits', function () {
+    $path = 'test.php';
+    $initialContent = "<?php\n\nnamespace App\\Models;\n\nclass User extends Model\n{\n    public function getName()\n    {\n        return \$this->name;\n    }\n}";
+    Storage::put($path, $initialContent);
+
+    (new File)->addTraits($path, 'HasFactory');
+
+    expect(Storage::get($path))->toBe("<?php\n\nnamespace App\\Models;\n\nclass User extends Model\n{\n    use HasFactory;\n\n    public function getName()\n    {\n        return \$this->name;\n    }\n}");
+});
+
+test('file->addTraits(...) adds trait to class with existing traits', function () {
+    $path = 'test.php';
+    $initialContent = "<?php\n\nnamespace App\\Models;\n\nclass User extends Model\n{\n    use HasFactory;\n\n    public function getName()\n    {\n        return \$this->name;\n    }\n}";
+    Storage::put($path, $initialContent);
+
+    (new File)->addTraits($path, 'Notifiable');
+
+    expect(Storage::get($path))->toBe("<?php\n\nnamespace App\\Models;\n\nclass User extends Model\n{\n    use HasFactory, Notifiable;\n\n    public function getName()\n    {\n        return \$this->name;\n    }\n}");
+});
+
+test('file->addTraits(...) is idempotent - does not add duplicate traits', function () {
+    $path = 'test.php';
+    $initialContent = "<?php\n\nnamespace App\\Models;\n\nclass User extends Model\n{\n    use HasFactory, Notifiable;\n\n    public function getName()\n    {\n        return \$this->name;\n    }\n}";
+    Storage::put($path, $initialContent);
+
+    (new File)->addTraits($path, 'HasFactory');
+
+    expect(Storage::get($path))->toBe($initialContent);
+});
+
+test('file->addTraits(...) adds multiple traits at once', function () {
+    $path = 'test.php';
+    $initialContent = "<?php\n\nnamespace App\\Models;\n\nclass User extends Model\n{\n    public function getName()\n    {\n        return \$this->name;\n    }\n}";
+    Storage::put($path, $initialContent);
+
+    (new File)->addTraits($path, ['HasFactory', 'Notifiable', 'HasUuids']);
+
+    expect(Storage::get($path))->toBe("<?php\n\nnamespace App\\Models;\n\nclass User extends Model\n{\n    use HasFactory, Notifiable, HasUuids;\n\n    public function getName()\n    {\n        return \$this->name;\n    }\n}");
+});
+
+test('file->addTraits(...) adds multiple traits to class with existing traits', function () {
+    $path = 'test.php';
+    $initialContent = "<?php\n\nnamespace App\\Models;\n\nclass User extends Model\n{\n    use HasFactory;\n\n    public function getName()\n    {\n        return \$this->name;\n    }\n}";
+    Storage::put($path, $initialContent);
+
+    (new File)->addTraits($path, ['Notifiable', 'HasUuids']);
+
+    expect(Storage::get($path))->toBe("<?php\n\nnamespace App\\Models;\n\nclass User extends Model\n{\n    use HasFactory, Notifiable, HasUuids;\n\n    public function getName()\n    {\n        return \$this->name;\n    }\n}");
+});
+
+test('file->addTraits(...) only adds traits not already present when adding multiple', function () {
+    $path = 'test.php';
+    $initialContent = "<?php\n\nnamespace App\\Models;\n\nclass User extends Model\n{\n    use HasFactory, Notifiable;\n\n    public function getName()\n    {\n        return \$this->name;\n    }\n}";
+    Storage::put($path, $initialContent);
+
+    (new File)->addTraits($path, ['HasFactory', 'HasUuids', 'Notifiable', 'SoftDeletes']);
+
+    expect(Storage::get($path))->toBe("<?php\n\nnamespace App\\Models;\n\nclass User extends Model\n{\n    use HasFactory, Notifiable, HasUuids, SoftDeletes;\n\n    public function getName()\n    {\n        return \$this->name;\n    }\n}");
+});
+
+test('file->addTraits(...) works with fully qualified trait names', function () {
+    $path = 'test.php';
+    $initialContent = "<?php\n\nnamespace App\\Models;\n\nclass User extends Model\n{\n    public function getName()\n    {\n        return \$this->name;\n    }\n}";
+    Storage::put($path, $initialContent);
+
+    (new File)->addTraits($path, 'Illuminate\\Database\\Eloquent\\Factories\\HasFactory');
+
+    expect(Storage::get($path))->toBe("<?php\n\nnamespace App\\Models;\n\nclass User extends Model\n{\n    use Illuminate\\Database\\Eloquent\\Factories\\HasFactory;\n\n    public function getName()\n    {\n        return \$this->name;\n    }\n}");
+});
+
+test('file->addTraits(...) throws exception when no class is found', function () {
+    $path = 'test.php';
+    $initialContent = "<?php\n\n// Just some code without a class";
+    Storage::put($path, $initialContent);
+
+    expect(fn () => (new File)->addTraits($path, 'HasFactory'))
+        ->toThrow(Exception::class, "Class not found in {$path}");
+});
+
+test('file->addTraits(...) preserves PHPDoc comments on traits', function () {
+    $path = 'test.php';
+    $initialContent = "<?php\n\nnamespace App\\Models;\n\nclass User extends Model\n{\n    /** @use HasFactory<\\Database\\Factories\\UserFactory> */\n    use HasFactory;\n\n    public function getName()\n    {\n        return \$this->name;\n    }\n}";
+    Storage::put($path, $initialContent);
+
+    (new File)->addTraits($path, 'Notifiable');
+
+    $result = Storage::get($path);
+    expect($result)->toContain('/** @use HasFactory<\\Database\\Factories\\UserFactory> */')
+        ->toContain('use HasFactory, Notifiable;');
+});