Added support for Bit objects with Psycopg 2 and pg8000

Author: astroknot-sheep

pgvector/pg8000/register.py

@@ -1,6 +1,6 @@
 from pg8000.native import Connection
 from typing import cast
-from .. import Vector, HalfVector, SparseVector
+from .. import Bit, Vector, HalfVector, SparseVector
 
 
 def register_vector(conn: Connection, /) -> None:
@@ -14,6 +14,8 @@ def register_vector(conn: Connection, /) -> None:
     conn.register_out_adapter(Vector, lambda v: v.to_text())
     conn.register_in_adapter(type_info['vector'], Vector.from_text)
 
+    conn.register_out_adapter(Bit, lambda v: v.to_text())
+
     try:
         import numpy as np
         conn.register_out_adapter(np.ndarray, lambda v: Vector(v).to_text())

pgvector/psycopg2/bit.py

@@ -0,0 +1,15 @@
+from psycopg2.extensions import adapt, register_adapter
+from typing import Any
+from .. import Bit
+
+
+class BitAdapter:
+    def __init__(self, value: Bit) -> None:
+        self._value = value
+
+    def getquoted(self) -> Any:
+        return adapt(self._value.to_text()).getquoted()
+
+
+def register_bit_info() -> None:
+    register_adapter(Bit, BitAdapter)

pgvector/psycopg2/register.py

@@ -1,5 +1,6 @@
 import psycopg2
 from psycopg2.extensions import connection, cursor
+from .bit import register_bit_info
 from .halfvec import register_halfvec_info
 from .sparsevec import register_sparsevec_info
 from .vector import register_vector_info
@@ -20,6 +21,8 @@ def register_vector(conn_or_curs: connection | cursor, /, *, globally: bool = Fa
 
     register_vector_info(type_info['vector'], type_info['_vector'] if arrays else None, scope)
 
+    register_bit_info()
+
     if 'halfvec' in type_info:
         register_halfvec_info(type_info['halfvec'], type_info['_halfvec'] if arrays else None, scope)
 

tests/test_pg8000.py

@@ -1,5 +1,5 @@
 from getpass import getuser
-from pgvector import HalfVector, SparseVector, Vector
+from pgvector import Bit, HalfVector, SparseVector, Vector
 from pgvector.pg8000 import register_vector
 from pg8000.native import Connection
 from .conftest import numpy as np
@@ -42,6 +42,14 @@ def test_bit(self) -> None:
         res = conn.run('SELECT binary_embedding FROM pg8000_items ORDER BY id')
         assert res == [['101'], [None]]
 
+    def test_bit_class(self) -> None:
+        embedding = Bit([True, False, True])
+        embedding2 = None
+        conn.run('INSERT INTO pg8000_items (binary_embedding) VALUES (:embedding), (:embedding2)', embedding=embedding, embedding2=embedding2)
+
+        res = conn.run('SELECT binary_embedding FROM pg8000_items ORDER BY id')
+        assert res == [['101'], [None]]
+
     def test_sparsevec(self) -> None:
         embedding = SparseVector([1.5, 2, 3])
         embedding2 = None

tests/test_psycopg2.py

@@ -1,4 +1,4 @@
-from pgvector import HalfVector, SparseVector, Vector
+from pgvector import Bit, HalfVector, SparseVector, Vector
 from pgvector.psycopg2 import register_vector
 import psycopg2
 from psycopg2.extras import DictCursor, RealDictCursor, NamedTupleCursor
@@ -61,6 +61,15 @@ def test_bit(self) -> None:
         assert res[0][0] == embedding
         assert res[1][0] is None
 
+    def test_bit_class(self) -> None:
+        embedding = Bit([True, False, True])
+        cur.execute('INSERT INTO psycopg2_items (binary_embedding) VALUES (%s), (NULL)', (embedding,))
+
+        cur.execute('SELECT binary_embedding FROM psycopg2_items ORDER BY id')
+        res = cur.fetchall()
+        assert res[0][0] == '101'
+        assert res[1][0] is None
+
     def test_sparsevec(self) -> None:
         embedding = SparseVector([1.5, 2, 3])
         cur.execute('INSERT INTO psycopg2_items (sparse_embedding) VALUES (%s), (NULL)', (embedding,))