Changed Distance to enum

pgvector · Feb 10, 2024 · 6ef7546 · 6ef7546
1 parent df8a547
commit 6ef7546
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 9 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,7 @@
+## 0.2.0 (unreleased)
+
+- Changed `Distance` to enum
+
 ## 0.1.4 (2023-11-14)
 
 - Moved package to `pgvector` namespace

diff --git a/src/laravel/Distance.php b/src/laravel/Distance.php
@@ -2,10 +2,9 @@
 
 namespace Pgvector\Laravel;
 
-// TODO use enum when PHP 8.0 reaches EOL
-class Distance
+enum Distance
 {
-    public const L2 = 0;
-    public const InnerProduct = 1;
-    public const Cosine = 2;
+    case L2;
+    case InnerProduct;
+    case Cosine;
 }
diff --git a/src/laravel/HasNeighbors.php b/src/laravel/HasNeighbors.php
@@ -6,7 +6,7 @@
 
 trait HasNeighbors
 {
-    public function scopeNearestNeighbors(Builder $query, string $column, mixed $value, int $distance): void
+    public function scopeNearestNeighbors(Builder $query, string $column, mixed $value, Distance $distance): void
     {
         switch ($distance) {
             case Distance::L2:
@@ -34,7 +34,7 @@ public function scopeNearestNeighbors(Builder $query, string $column, mixed $val
             ->orderByRaw($order, [$vector]);
     }
 
-    public function nearestNeighbors(string $column, int $distance): Builder
+    public function nearestNeighbors(string $column, Distance $distance): Builder
     {
         $id = $this->getKey();
         if (!array_key_exists($column, $this->attributes)) {

diff --git a/tests/LaravelTest.php b/tests/LaravelTest.php
@@ -108,8 +108,7 @@ public function testMissingAttribute()
 
     public function testInvalidDistance()
     {
-        $this->expectException(InvalidArgumentException::class);
-        $this->expectExceptionMessage('Invalid distance');
+        $this->expectException(TypeError::class);
 
         Item::query()->nearestNeighbors('embedding', [1, 2, 3], 4);
     }