Skip to content

Commit c223ef4

Browse files
ankaneankane
committed
Improved optional test imports
1 parent ad4a7e8 commit c223ef4

10 files changed

Lines changed: 102 additions & 98 deletions

‎tests/conftest.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
try:
2+
import numpy
3+
except ImportError:
4+
numpy = None # type: ignore
5+
6+
__all__ = [
7+
'numpy'
8+
]

‎tests/test_asyncpg.py‎

Lines changed: 4 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,7 @@
33
from pgvector import HalfVector, SparseVector, Vector
44
from pgvector.asyncpg import register_vector
55
import pytest
6-
7-
try:
8-
import numpy as np
9-
NUMPY_AVAILABLE = True
10-
except ImportError:
11-
NUMPY_AVAILABLE = False
6+
from .conftest import numpy as np
127

138

149
class TestAsyncpg:
@@ -26,7 +21,7 @@ async def test_vector(self) -> None:
2621

2722
embedding = Vector([1.5, 2, 3])
2823
embedding2 = [4.5, 5, 6]
29-
embedding3 = np.array([7.5, 8, 9]) if NUMPY_AVAILABLE else [7.5, 8, 9]
24+
embedding3 = np.array([7.5, 8, 9]) if np is not None else [7.5, 8, 9]
3025
embedding4 = None
3126
await conn.execute('INSERT INTO asyncpg_items (embedding) VALUES ($1), ($2), ($3), ($4)', embedding, embedding2, embedding3, embedding4)
3227

@@ -117,14 +112,14 @@ async def test_vector_array(self) -> None:
117112
embeddings2 = [[1.5, 2, 3], [4.5, 5, 6]]
118113
await conn.execute('INSERT INTO asyncpg_items (embeddings) VALUES (ARRAY[$1, $2]::vector[])', embeddings2[0], embeddings2[1])
119114

120-
if NUMPY_AVAILABLE:
115+
if np is not None:
121116
embeddings3 = [np.array([1.5, 2, 3]), np.array([4.5, 5, 6])]
122117
await conn.execute('INSERT INTO asyncpg_items (embeddings) VALUES (ARRAY[$1, $2]::vector[])', embeddings3[0], embeddings3[1])
123118

124119
res = await conn.fetch('SELECT * FROM asyncpg_items ORDER BY id')
125120
assert res[0]['embeddings'] == embeddings
126121
assert res[1]['embeddings'] == [Vector(e) for e in embeddings2]
127-
if NUMPY_AVAILABLE:
122+
if np is not None:
128123
assert res[2]['embeddings'] == [Vector(e) for e in embeddings3] # type: ignore
129124

130125
await conn.close()

‎tests/test_bit.py‎

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,7 @@
22
import pytest
33
import random
44
from struct import pack
5-
6-
try:
7-
import numpy as np
8-
NUMPY_AVAILABLE = True
9-
except ImportError:
10-
NUMPY_AVAILABLE = False
5+
from .conftest import numpy as np
116

127

138
class TestBit:
@@ -41,26 +36,34 @@ def test_bytes(self) -> None:
4136
assert Bit(b'\xff\x00\xf0').to_text() == '111111110000000011110000'
4237
assert Bit(b'\xfe\x07\x00').to_text() == '111111100000011100000000'
4338

44-
@pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required')
4539
def test_ndarray(self) -> None:
40+
if np is None:
41+
pytest.skip('NumPy required')
42+
4643
arr = np.array([True, False, True])
4744
assert Bit(arr).to_list() == [True, False, True]
4845
assert np.array_equal(Bit(arr).to_numpy(), arr)
4946

50-
@pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required')
5147
def test_ndarray_unpackbits(self) -> None:
48+
if np is None:
49+
pytest.skip('NumPy required')
50+
5251
arr = np.unpackbits(np.array([254, 7, 0], dtype=np.uint8))
5352
assert Bit(arr).to_text() == '111111100000011100000000'
5453

55-
@pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required')
5654
def test_ndarray_uint8(self) -> None:
55+
if np is None:
56+
pytest.skip('NumPy required')
57+
5758
arr = np.array([254, 7, 0], dtype=np.uint8)
5859
with pytest.raises(ValueError) as error:
5960
Bit(arr)
6061
assert str(error.value) == 'expected elements to be boolean'
6162

62-
@pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required')
6363
def test_ndarray_uint16(self) -> None:
64+
if np is None:
65+
pytest.skip('NumPy required')
66+
6467
arr = np.array([254, 7, 0], dtype=np.uint16)
6568
with pytest.raises(ValueError) as error:
6669
Bit(arr) # type: ignore
@@ -86,14 +89,14 @@ def test_equality(self) -> None:
8689
def test_from_text(self) -> None:
8790
vec = Bit.from_text('101')
8891
assert vec.to_list() == [True, False, True]
89-
if NUMPY_AVAILABLE:
92+
if np is not None:
9093
assert np.array_equal(vec.to_numpy(), [True, False, True])
9194
assert vec.to_text() == '101'
9295

9396
def test_from_binary(self) -> None:
9497
data = pack('>iB', 3, 5 << 5)
9598
vec = Bit.from_binary(data)
9699
assert vec.to_list() == [True, False, True]
97-
if NUMPY_AVAILABLE:
100+
if np is not None:
98101
assert np.array_equal(vec.to_numpy(), [True, False, True])
99102
assert vec.to_binary() == data

‎tests/test_half_vector.py‎

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,7 @@
11
from pgvector import HalfVector
22
import pytest
33
from struct import pack
4-
5-
try:
6-
import numpy as np
7-
NUMPY_AVAILABLE = True
8-
except ImportError:
9-
NUMPY_AVAILABLE = False
4+
from .conftest import numpy as np
105

116

127
class TestHalfVector:
@@ -28,8 +23,10 @@ def test_list_list(self) -> None:
2823
HalfVector([[1, 2], [3, 4]]) # type: ignore
2924
assert str(error.value) == 'expected list[float]'
3025

31-
@pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required')
3226
def test_ndarray(self) -> None:
27+
if np is None:
28+
pytest.skip('NumPy required')
29+
3330
arr = np.array([1, 2, 3], dtype=np.float16)
3431
assert HalfVector(arr).to_list() == [1, 2, 3]
3532
assert HalfVector(arr).to_numpy() is not arr
@@ -64,14 +61,14 @@ def test_dimensions(self) -> None:
6461
def test_from_text(self) -> None:
6562
vec = HalfVector.from_text('[1.5,2,3]')
6663
assert vec.to_list() == [1.5, 2, 3]
67-
if NUMPY_AVAILABLE:
64+
if np is not None:
6865
assert np.array_equal(vec.to_numpy(), [1.5, 2, 3])
6966
assert vec.to_text() == '[1.5,2.0,3.0]'
7067

7168
def test_from_binary(self) -> None:
7269
data = pack('>HH3e', 3, 0, 1.5, 2, 3)
7370
vec = HalfVector.from_binary(data)
7471
assert vec.to_list() == [1.5, 2, 3]
75-
if NUMPY_AVAILABLE:
72+
if np is not None:
7673
assert np.array_equal(vec.to_numpy(), [1.5, 2, 3])
7774
assert vec.to_binary() == data

‎tests/test_pg8000.py‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,7 @@
22
from pgvector import HalfVector, SparseVector, Vector
33
from pgvector.pg8000 import register_vector
44
from pg8000.native import Connection
5-
6-
try:
7-
import numpy as np
8-
NUMPY_AVAILABLE = True
9-
except ImportError:
10-
NUMPY_AVAILABLE = False
5+
from .conftest import numpy as np
116

127
conn = Connection(getuser(), database='pgvector_python_test')
138

@@ -24,7 +19,7 @@ def setup_method(self) -> None:
2419

2520
def test_vector(self) -> None:
2621
embedding = Vector([1.5, 2, 3])
27-
embedding2 = np.array([4.5, 5, 6]) if NUMPY_AVAILABLE else Vector([4.5, 5, 6])
22+
embedding2 = np.array([4.5, 5, 6]) if np is not None else Vector([4.5, 5, 6])
2823
embedding3 = None
2924
conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (:embedding2), (:embedding3)', embedding=embedding, embedding2=embedding2, embedding3=embedding3)
3025

‎tests/test_psycopg.py‎

Lines changed: 13 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,7 @@
55
from psycopg_pool import ConnectionPool, AsyncConnectionPool
66
import pytest
77
from typing import Any
8-
9-
try:
10-
import numpy as np
11-
NUMPY_AVAILABLE = True
12-
except ImportError:
13-
NUMPY_AVAILABLE = False
8+
from .conftest import numpy as np
149

1510
conn = psycopg.connect(dbname='pgvector_python_test', autocommit=True)
1611

@@ -49,27 +44,35 @@ def test_vector_binary_format_correct(self) -> None:
4944
res = next(conn.execute('SELECT %b::vector::text', (embedding,)))[0]
5045
assert res == '[1.5,2,3]'
5146

52-
@pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required')
5347
def test_vector_numpy_binary_format(self) -> None:
48+
if np is None:
49+
pytest.skip('NumPy required')
50+
5451
embedding = np.array([1.5, 2, 3])
5552
res = next(conn.execute('SELECT %b::vector', (embedding,), binary=True))[0]
5653
assert res == Vector(embedding)
5754

58-
@pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required')
5955
def test_vector_numpy_text_format(self) -> None:
56+
if np is None:
57+
pytest.skip('NumPy required')
58+
6059
embedding = np.array([1.5, 2, 3])
6160
res = next(conn.execute('SELECT %t::vector', (embedding,)))[0]
6261
assert res == Vector(embedding)
6362

64-
@pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required')
6563
def test_vector_numpy_binary_format_non_contiguous(self) -> None:
64+
if np is None:
65+
pytest.skip('NumPy required')
66+
6667
embedding = np.flipud(np.array([1.5, 2, 3]))
6768
assert not embedding.data.contiguous
6869
res = next(conn.execute('SELECT %b::vector', (embedding,)))[0]
6970
assert res == Vector([3, 2, 1.5])
7071

71-
@pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required')
7272
def test_vector_numpy_text_format_non_contiguous(self) -> None:
73+
if np is None:
74+
pytest.skip('NumPy required')
75+
7376
embedding = np.flipud(np.array([1.5, 2, 3]))
7477
assert not embedding.data.contiguous
7578
res = next(conn.execute('SELECT %t::vector', (embedding,)))[0]

‎tests/test_psycopg2.py‎

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,7 @@
44
from psycopg2.extras import DictCursor, RealDictCursor, NamedTupleCursor
55
from psycopg2.pool import ThreadedConnectionPool
66
import pytest
7-
8-
try:
9-
import numpy as np
10-
NUMPY_AVAILABLE = True
11-
except ImportError:
12-
NUMPY_AVAILABLE = False
7+
from .conftest import numpy as np
138

149
conn = psycopg2.connect(dbname='pgvector_python_test')
1510
conn.autocommit = True
@@ -36,8 +31,10 @@ def test_vector(self) -> None:
3631
assert res[0][0] == embedding
3732
assert res[1][0] is None
3833

39-
@pytest.mark.skipif(not NUMPY_AVAILABLE, reason='NumPy required')
4034
def test_vector_numpy(self) -> None:
35+
if np is None:
36+
pytest.skip('NumPy required')
37+
4138
embedding = np.array([1.5, 2, 3])
4239
cur.execute('INSERT INTO psycopg2_items (embedding) VALUES (%s), (NULL)', (embedding,))
4340

0 commit comments

Comments
 (0)