Skip to content

Commit f610d5e

Browse files
ankaneankane
committed
Fixed typechecking [skip ci]
1 parent 2bf00f8 commit f610d5e

14 files changed

Lines changed: 56 additions & 42 deletions

‎pgvector/bit.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,9 @@
44

55
try:
66
import numpy as np
7+
NUMPY_AVAILABLE = True
78
except ImportError:
8-
np = None
9+
NUMPY_AVAILABLE = False
910

1011

1112
class Bit:
@@ -34,7 +35,7 @@ def bit_value(v: bool) -> str:
3435
data = int(value, 2).to_bytes(len(value) // 8, byteorder='big')
3536
except ValueError:
3637
raise ValueError('expected bit string')
37-
elif np is not None and isinstance(value, np.ndarray):
38+
elif NUMPY_AVAILABLE and isinstance(value, np.ndarray):
3839
if value.dtype != np.bool:
3940
# skip warning for result of np.unpackbits
4041
if value.dtype != np.uint8 or np.any(value > 1):

‎pgvector/halfvec.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,9 @@
33

44
try:
55
import numpy as np
6+
NUMPY_AVAILABLE = True
67
except ImportError:
7-
np = None
8+
NUMPY_AVAILABLE = False
89

910

1011
class HalfVector:
@@ -15,7 +16,7 @@ def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.float
1516
self._value = struct.pack(f'>HH{dim}e', dim, 0, *value)
1617
except struct.error:
1718
raise ValueError('expected list[float]')
18-
elif np is not None and isinstance(value, np.ndarray):
19+
elif NUMPY_AVAILABLE and isinstance(value, np.ndarray):
1920
if value.ndim != 1:
2021
raise ValueError('expected ndim to be 1')
2122

‎pgvector/sparsevec.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
try:
66
import numpy as np
77
except ImportError:
8-
np = None
8+
pass
99

1010
NO_DEFAULT = object()
1111

‎pgvector/vector.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,9 @@
33

44
try:
55
import numpy as np
6+
NUMPY_AVAILABLE = True
67
except ImportError:
7-
np = None
8+
NUMPY_AVAILABLE = False
89

910

1011
class Vector:
@@ -15,7 +16,7 @@ def __init__(self, value: list[float] | np.ndarray[tuple[int], np.dtype[np.float
1516
self._value = struct.pack(f'>HH{dim}f', dim, 0, *value)
1617
except struct.error:
1718
raise ValueError('expected list[float]')
18-
elif np is not None and isinstance(value, np.ndarray):
19+
elif NUMPY_AVAILABLE and isinstance(value, np.ndarray):
1920
if value.ndim != 1:
2021
raise ValueError('expected ndim to be 1')
2122

‎tests/test_bit.py‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,9 @@
33

44
try:
55
import numpy as np
6+
NUMPY_AVAILABLE = True
67
except ImportError:
7-
np = None
8+
NUMPY_AVAILABLE = False
89

910

1011
class TestBit:
@@ -33,24 +34,24 @@ def test_bytes(self):
3334
assert Bit(b'\xff\x00\xf0').to_text() == '111111110000000011110000'
3435
assert Bit(b'\xfe\x07\x00').to_text() == '111111100000011100000000'
3536

36-
@pytest.mark.skipif(np is None, reason='NumPy required')
37+
@pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required')
3738
def test_ndarray(self):
3839
arr = np.array([True, False, True])
3940
assert Bit(arr).to_list() == [True, False, True]
4041
assert np.array_equal(Bit(arr).to_numpy(), arr)
4142

42-
@pytest.mark.skipif(np is None, reason='NumPy required')
43+
@pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required')
4344
def test_ndarray_unpackbits(self):
4445
arr = np.unpackbits(np.array([254, 7, 0], dtype=np.uint8))
4546
assert Bit(arr).to_text() == '111111100000011100000000'
4647

47-
@pytest.mark.skipif(np is None, reason='NumPy required')
48+
@pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required')
4849
def test_ndarray_uint8(self):
4950
arr = np.array([254, 7, 0], dtype=np.uint8)
5051
with pytest.warns(UserWarning, match='expected elements to be boolean'):
5152
assert Bit(arr).to_text() == '110'
5253

53-
@pytest.mark.skipif(np is None, reason='NumPy required')
54+
@pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required')
5455
def test_ndarray_uint16(self):
5556
arr = np.array([254, 7, 0], dtype=np.uint16)
5657
with pytest.warns(UserWarning, match='expected elements to be boolean'):

‎tests/test_django.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,9 @@
1919

2020
try:
2121
import numpy as np
22+
NUMPY_AVAILABLE = True
2223
except ImportError:
23-
np = None
24+
NUMPY_AVAILABLE = False
2425

2526
settings.configure(
2627
DATABASES={

‎tests/test_half_vector.py‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,9 @@
44

55
try:
66
import numpy as np
7+
NUMPY_AVAILABLE = True
78
except ImportError:
8-
np = None
9+
NUMPY_AVAILABLE = False
910

1011

1112
class TestHalfVector:
@@ -25,7 +26,7 @@ def test_list_list(self):
2526
HalfVector([[1, 2], [3, 4]]) # ty: ignore[invalid-argument-type]
2627
assert str(error.value) == 'expected list[float]'
2728

28-
@pytest.mark.skipif(np is None, reason='NumPy required')
29+
@pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required')
2930
def test_ndarray(self):
3031
arr = np.array([1, 2, 3])
3132
assert HalfVector(arr).to_list() == [1, 2, 3]
@@ -47,7 +48,7 @@ def test_equality(self):
4748
def test_dimensions(self):
4849
assert HalfVector([1, 2, 3]).dimensions() == 3
4950

50-
@pytest.mark.skipif(np is None, reason='NumPy required')
51+
@pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required')
5152
def test_to_numpy_readonly(self):
5253
arr = HalfVector([1, 2, 3]).to_numpy()
5354
with pytest.raises(ValueError) as error:

‎tests/test_pg8000.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,9 @@
66

77
try:
88
import numpy as np
9+
NUMPY_AVAILABLE = True
910
except ImportError:
10-
np = None
11+
NUMPY_AVAILABLE = False
1112

1213
conn = Connection(getuser(), database='pgvector_python_test')
1314

@@ -30,7 +31,7 @@ def test_vector(self):
3031
assert res[0][0] == embedding
3132
assert res[1][0] is None
3233

33-
@pytest.mark.skipif(np is None, reason='NumPy required')
34+
@pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required')
3435
def test_vector_numpy(self):
3536
embedding = np.array([1.5, 2, 3])
3637
conn.run('INSERT INTO pg8000_items (embedding) VALUES (:embedding), (NULL)', embedding=embedding)

‎tests/test_psycopg.py‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,9 @@
66

77
try:
88
import numpy as np
9+
NUMPY_AVAILABLE = True
910
except ImportError:
10-
np = None
11+
NUMPY_AVAILABLE = False
1112

1213
conn = psycopg.connect(dbname='pgvector_python_test', autocommit=True)
1314

@@ -45,26 +46,26 @@ def test_vector_binary_format_correct(self):
4546
res = next(conn.execute('SELECT %b::vector::text', (embedding,)))[0]
4647
assert res == '[1.5,2,3]'
4748

48-
@pytest.mark.skipif(np is None, reason='NumPy required')
49+
@pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required')
4950
def test_vector_numpy_binary_format(self):
5051
embedding = np.array([1.5, 2, 3])
5152
res = next(conn.execute('SELECT %b::vector', (embedding,), binary=True))[0]
5253
assert res == Vector(embedding)
5354

54-
@pytest.mark.skipif(np is None, reason='NumPy required')
55+
@pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required')
5556
def test_vector_numpy_text_format(self):
5657
embedding = np.array([1.5, 2, 3])
5758
res = next(conn.execute('SELECT %t::vector', (embedding,)))[0]
5859
assert res == Vector(embedding)
5960

60-
@pytest.mark.skipif(np is None, reason='NumPy required')
61+
@pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required')
6162
def test_vector_numpy_binary_format_non_contiguous(self):
6263
embedding = np.flipud(np.array([1.5, 2, 3]))
6364
assert not embedding.data.contiguous
6465
res = next(conn.execute('SELECT %b::vector', (embedding,)))[0]
6566
assert res == Vector([3, 2, 1.5])
6667

67-
@pytest.mark.skipif(np is None, reason='NumPy required')
68+
@pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required')
6869
def test_vector_numpy_text_format_non_contiguous(self):
6970
embedding = np.flipud(np.array([1.5, 2, 3]))
7071
assert not embedding.data.contiguous

‎tests/test_psycopg2.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,9 @@
77

88
try:
99
import numpy as np
10+
NUMPY_AVAILABLE = True
1011
except ImportError:
11-
np = None
12+
NUMPY_AVAILABLE = False
1213

1314
conn = psycopg2.connect(dbname='pgvector_python_test')
1415
conn.autocommit = True
@@ -34,7 +35,7 @@ def test_vector(self):
3435
assert res[0][0] == embedding
3536
assert res[1][0] is None
3637

37-
@pytest.mark.skipif(np is None, reason='NumPy required')
38+
@pytest.mark.skipif(NUMPY_AVAILABLE, reason='NumPy required')
3839
def test_vector_numpy(self):
3940
embedding = np.array([1.5, 2, 3])
4041
cur.execute('INSERT INTO psycopg2_items (embedding) VALUES (%s), (NULL)', (embedding,))

0 commit comments

Comments
 (0)