Switched to TypeDecorator for bit type for SQLAlchemy

Author: ankane

pgvector/sqlalchemy/bit.py

@@ -1,42 +1,24 @@
-from sqlalchemy.dialects.postgresql.base import ischema_names
-from sqlalchemy.types import UserDefinedType, TypeEngine, Float
-from sqlalchemy import Dialect, Operators
+from sqlalchemy.dialects.postgresql.base import PGBit
+from sqlalchemy.types import TypeDecorator, Float
+from sqlalchemy import Operators
 from typing import Any
 
 
-class BIT(UserDefinedType[Any]):
+class BIT(TypeDecorator[Any]):
+    impl = PGBit
     cache_ok = True
 
-    def __init__(self, length: int | None = None) -> None:
-        super().__init__()
-        self.length = length
-
-    def get_col_spec(self, **kw: Any) -> str:
-        if self.length is None:
-            return 'BIT'
-        return 'BIT(%d)' % self.length
-
-    def bind_processor(self, dialect: Dialect) -> Any:
-        if dialect.__class__.__name__ == 'PGDialect_asyncpg':
+    def process_bind_param(self, value: Any, dialect: Any) -> Any:
+        if dialect.__class__.__name__ == 'PGDialect_asyncpg' and isinstance(value, str):
             import asyncpg
+            return asyncpg.BitString(value)  # type: ignore
+        return value
 
-            def process(value: Any) -> Any:
-                if isinstance(value, str):
-                    return asyncpg.BitString(value)  # type: ignore
-                return value
-            return process
-        else:
-            return super().bind_processor(dialect)
-
-    class Comparator(TypeEngine.Comparator[Any]):
+    class Comparator(TypeDecorator.Comparator[Any]):
         def hamming_distance(self, other: object) -> Operators:
             return self.op('<~>', return_type=Float)(other)
 
         def jaccard_distance(self, other: object) -> Operators:
             return self.op('<%>', return_type=Float)(other)
 
-    comparator_factory = Comparator
-
-
-# for reflection
-ischema_names['bit'] = BIT  # type: ignore
+    comparator_factory = Comparator  # type: ignore