Skip to content

Commit ad4a7e8

Browse files
ankaneankane
committed
Removed unnecessary importing
1 parent 24a95cb commit ad4a7e8

7 files changed

Lines changed: 54 additions & 40 deletions

File tree

‎pgvector/_utils.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
import sys
2+
3+
4+
def is_ndarray(value: object) -> bool:
5+
if (numpy := sys.modules.get('numpy')):
6+
return isinstance(value, numpy.ndarray)
7+
return False
8+
9+
10+
def is_sparse_array(value: object) -> bool:
11+
if (sparse := sys.modules.get('scipy.sparse')):
12+
return isinstance(value, (sparse.sparray, sparse.spmatrix))
13+
return False

‎pgvector/bit.py‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,10 @@
11
from __future__ import annotations
22
from struct import pack, unpack_from
3+
from typing import TYPE_CHECKING
4+
from ._utils import is_ndarray
35

4-
try:
6+
if TYPE_CHECKING:
57
import numpy as np
6-
NUMPY_AVAILABLE = True
7-
except ImportError:
8-
NUMPY_AVAILABLE = False
98

109

1110
class Bit:
@@ -33,7 +32,9 @@ def bit_value(v: bool) -> str:
3332
self._data = int(value, 2).to_bytes(len(value) // 8, byteorder='big')
3433
except ValueError:
3534
raise ValueError('expected bit string')
36-
elif NUMPY_AVAILABLE and isinstance(value, np.ndarray):
35+
elif is_ndarray(value):
36+
import numpy as np
37+
3738
if value.dtype != np.bool:
3839
# skip error for result of np.unpackbits
3940
if value.dtype != np.uint8 or np.any(value > 1):
@@ -61,6 +62,8 @@ def to_list(self) -> list[bool]:
6162
return [v != '0' for v in self.to_text()]
6263

6364
def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.bool]]:
65+
import numpy as np
66+
6467
return np.unpackbits(np.frombuffer(self._data, dtype=np.uint8), count=self._length).astype(bool)
6568

6669
def to_text(self) -> str:

‎pgvector/halfvec.py‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,11 @@
22
import array
33
import struct
44
import sys
5+
from typing import TYPE_CHECKING
6+
from ._utils import is_ndarray
57

6-
try:
8+
if TYPE_CHECKING:
79
import numpy as np
8-
NUMPY_AVAILABLE = True
9-
except ImportError:
10-
NUMPY_AVAILABLE = False
1110

1211

1312
class HalfVector:
@@ -18,7 +17,9 @@ def __init__(self, value: list[float] | np.ndarray[tuple[int, ...], np.dtype[np.
1817
self._value = array.array('H', struct.pack(f'{dim}e', *value))
1918
except struct.error:
2019
raise ValueError('expected list[float]')
21-
elif NUMPY_AVAILABLE and isinstance(value, np.ndarray):
20+
elif is_ndarray(value):
21+
import numpy as np
22+
2223
if value.ndim != 1:
2324
raise ValueError('expected ndim to be 1')
2425

@@ -44,6 +45,7 @@ def to_list(self) -> list[float]:
4445
return list(struct.unpack(f'{dim}e', self._value))
4546

4647
def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.float16]]:
48+
import numpy as np
4749
return np.frombuffer(self._value, dtype=np.float16)
4850

4951
def to_text(self) -> str:

‎pgvector/psycopg/vector.py‎

Lines changed: 4 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4,16 +4,13 @@
44
from psycopg.adapt import Loader, Dumper
55
from psycopg.pq import Format
66
from psycopg.types import TypeInfo
7-
from typing import Any, TypeAlias
7+
from typing import TYPE_CHECKING, Any, TypeAlias
88
from .. import Vector
99

1010
Buffer: TypeAlias = bytes | bytearray | memoryview
1111

12-
try:
12+
if TYPE_CHECKING:
1313
import numpy as np
14-
NUMPY_AVAILABLE = True
15-
except ImportError:
16-
NUMPY_AVAILABLE = False
1714

1815

1916
class VectorDumper(Dumper):
@@ -62,9 +59,7 @@ def register_vector_info(context: BaseConnection[Any], info: TypeInfo | None) ->
6259
adapters = context.adapters
6360
adapters.register_dumper(Vector, text_dumper)
6461
adapters.register_dumper(Vector, binary_dumper)
62+
adapters.register_dumper('numpy.ndarray', text_dumper)
63+
adapters.register_dumper('numpy.ndarray', binary_dumper)
6564
adapters.register_loader(info.oid, VectorLoader)
6665
adapters.register_loader(info.oid, VectorBinaryLoader)
67-
68-
if NUMPY_AVAILABLE:
69-
adapters.register_dumper(np.ndarray, text_dumper)
70-
adapters.register_dumper(np.ndarray, binary_dumper)

‎pgvector/psycopg2/vector.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,10 @@
11
from __future__ import annotations
22
from psycopg2.extensions import adapt, connection, cursor, new_array_type, new_type, register_adapter, register_type
3-
from typing import Any
3+
from typing import TYPE_CHECKING, Any
44
from .. import Vector
55

6-
try:
6+
if TYPE_CHECKING:
77
import numpy as np
8-
NUMPY_AVAILABLE = True
9-
except ImportError:
10-
NUMPY_AVAILABLE = False
118

129

1310
class VectorAdapter:
@@ -38,5 +35,8 @@ def register_vector_info(oid: int, array_oid: int | None, scope: connection | cu
3835

3936
register_adapter(Vector, VectorAdapter)
4037

41-
if NUMPY_AVAILABLE:
38+
try:
39+
import numpy as np
4240
register_adapter(np.ndarray, VectorAdapter)
41+
except ImportError:
42+
pass

‎pgvector/sparsevec.py‎

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,12 @@
11
from __future__ import annotations
22
from struct import pack, unpack_from
3-
from typing import Any, overload
3+
from typing import TYPE_CHECKING, Any, overload
4+
from ._utils import is_sparse_array
45

5-
try:
6+
if TYPE_CHECKING:
67
import numpy as np
7-
except ImportError:
8-
pass
9-
10-
try:
118
from scipy.sparse import sparray, spmatrix, coo_array, coo_matrix
12-
SCIPY_AVAILABLE = True
13-
except ImportError:
14-
SCIPY_AVAILABLE = False
9+
1510

1611
NO_DEFAULT = object()
1712

@@ -26,7 +21,7 @@ def __init__(self, value: list[float] | np.ndarray[tuple[int, ...], np.dtype[np.
2621
...
2722

2823
def __init__(self, value: dict[int, float] | list[float] | np.ndarray[tuple[int, ...], np.dtype[np.floating]] | sparray | spmatrix, dimensions: int | Any = NO_DEFAULT, /) -> None:
29-
if SCIPY_AVAILABLE and isinstance(value, (sparray, spmatrix)):
24+
if is_sparse_array(value):
3025
if dimensions is not NO_DEFAULT:
3126
raise ValueError('extra argument')
3227

@@ -61,6 +56,8 @@ def values(self) -> list[float]:
6156
return self._values
6257

6358
def to_coo(self) -> coo_array:
59+
from scipy.sparse import coo_array
60+
6461
coords = ([0] * len(self._indices), self._indices)
6562
return coo_array((self._values, coords), shape=(1, self._dim))
6663

@@ -71,6 +68,8 @@ def to_list(self) -> list[float]:
7168
return vec
7269

7370
def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.float32]]:
71+
import numpy as np
72+
7473
vec = np.repeat(0.0, self._dim).astype(np.float32)
7574
for i, v in zip(self._indices, self._values):
7675
vec[i] = v

‎pgvector/vector.py‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,11 @@
22
import array
33
import struct
44
import sys
5+
from typing import TYPE_CHECKING
6+
from ._utils import is_ndarray
57

6-
try:
8+
if TYPE_CHECKING:
79
import numpy as np
8-
NUMPY_AVAILABLE = True
9-
except ImportError:
10-
NUMPY_AVAILABLE = False
1110

1211

1312
class Vector:
@@ -17,7 +16,9 @@ def __init__(self, value: list[float] | np.ndarray[tuple[int, ...], np.dtype[np.
1716
self._value = array.array('f', value)
1817
except TypeError:
1918
raise ValueError('expected list[float]')
20-
elif NUMPY_AVAILABLE and isinstance(value, np.ndarray):
19+
elif is_ndarray(value):
20+
import numpy as np
21+
2122
if value.ndim != 1:
2223
raise ValueError('expected ndim to be 1')
2324

@@ -42,6 +43,7 @@ def to_list(self) -> list[float]:
4243
return self._value.tolist()
4344

4445
def to_numpy(self) -> np.ndarray[tuple[int, ...], np.dtype[np.float32]]:
46+
import numpy as np
4547
return np.frombuffer(self._value, dtype=np.float32)
4648

4749
def to_text(self) -> str:

0 commit comments

Comments
 (0)