Moved Vector to separate file for SQLAlchemy [skip ci]

ankane · ankane · commit afb416f947a2 · 2024-05-16T16:31:04.000-04:00
diff --git a/pgvector/sqlalchemy/__init__.py b/pgvector/sqlalchemy/__init__.py
@@ -1,53 +1,3 @@
-from sqlalchemy.dialects.postgresql.base import ischema_names
-from sqlalchemy.types import UserDefinedType, Float, String
-from ..utils import from_db, to_db
+from .vector import Vector
 
 __all__ = ['Vector']
-
-
-class Vector(UserDefinedType):
-    cache_ok = True
-    _string = String()
-
-    def __init__(self, dim=None):
-        super(UserDefinedType, self).__init__()
-        self.dim = dim
-
-    def get_col_spec(self, **kw):
-        if self.dim is None:
-            return 'VECTOR'
-        return 'VECTOR(%d)' % self.dim
-
-    def bind_processor(self, dialect):
-        def process(value):
-            return to_db(value, self.dim)
-        return process
-
-    def literal_processor(self, dialect):
-        string_literal_processor = self._string._cached_literal_processor(dialect)
-
-        def process(value):
-            return string_literal_processor(to_db(value, self.dim))
-        return process
-
-    def result_processor(self, dialect, coltype):
-        def process(value):
-            return from_db(value)
-        return process
-
-    class comparator_factory(UserDefinedType.Comparator):
-        def l2_distance(self, other):
-            return self.op('<->', return_type=Float)(other)
-
-        def max_inner_product(self, other):
-            return self.op('<#>', return_type=Float)(other)
-
-        def cosine_distance(self, other):
-            return self.op('<=>', return_type=Float)(other)
-
-        def l1_distance(self, other):
-            return self.op('<+>', return_type=Float)(other)
-
-
-# for reflection
-ischema_names['vector'] = Vector
diff --git a/pgvector/sqlalchemy/vector.py b/pgvector/sqlalchemy/vector.py
@@ -0,0 +1,51 @@
+from sqlalchemy.dialects.postgresql.base import ischema_names
+from sqlalchemy.types import UserDefinedType, Float, String
+from ..utils import from_db, to_db
+
+
+class Vector(UserDefinedType):
+    cache_ok = True
+    _string = String()
+
+    def __init__(self, dim=None):
+        super(UserDefinedType, self).__init__()
+        self.dim = dim
+
+    def get_col_spec(self, **kw):
+        if self.dim is None:
+            return 'VECTOR'
+        return 'VECTOR(%d)' % self.dim
+
+    def bind_processor(self, dialect):
+        def process(value):
+            return to_db(value, self.dim)
+        return process
+
+    def literal_processor(self, dialect):
+        string_literal_processor = self._string._cached_literal_processor(dialect)
+
+        def process(value):
+            return string_literal_processor(to_db(value, self.dim))
+        return process
+
+    def result_processor(self, dialect, coltype):
+        def process(value):
+            return from_db(value)
+        return process
+
+    class comparator_factory(UserDefinedType.Comparator):
+        def l2_distance(self, other):
+            return self.op('<->', return_type=Float)(other)
+
+        def max_inner_product(self, other):
+            return self.op('<#>', return_type=Float)(other)
+
+        def cosine_distance(self, other):
+            return self.op('<=>', return_type=Float)(other)
+
+        def l1_distance(self, other):
+            return self.op('<+>', return_type=Float)(other)
+
+
+# for reflection
+ischema_names['vector'] = Vector
