Improved ndarray imports

Author: ankane

pgvector/_utils.py

@@ -5,6 +5,10 @@
     import numpy as np
 
     ndarray: TypeAlias = np.ndarray[tuple[int, ...], np.dtype[np.floating]]
+else:
+    # any value works since not type checking
+    # TODO use Never when Python 3.10 no longer supported
+    ndarray = None
 
 
 def is_ndarray(value: object) -> bool:

pgvector/halfvec.py

@@ -3,11 +3,10 @@
 import struct
 import sys
 from typing import TYPE_CHECKING
-from ._utils import is_ndarray
+from ._utils import is_ndarray, ndarray
 
 if TYPE_CHECKING:
     import numpy as np
-    from ._utils import ndarray
 
 
 class HalfVector:

pgvector/psycopg/vector.py

@@ -5,11 +5,9 @@
 from psycopg.adapt import Loader, Dumper
 from psycopg.pq import Format
 from psycopg.types import TypeInfo
-from typing import TYPE_CHECKING, Any
+from typing import Any
 from .. import Vector
-
-if TYPE_CHECKING:
-    from .._utils import ndarray
+from .._utils import ndarray
 
 
 class VectorDumper(Dumper):

pgvector/psycopg2/vector.py

@@ -1,10 +1,8 @@
 from __future__ import annotations
 from psycopg2.extensions import adapt, connection, cursor, new_array_type, new_type, register_adapter, register_type
-from typing import TYPE_CHECKING, Any
+from typing import Any
 from .. import Vector
-
-if TYPE_CHECKING:
-    from .._utils import ndarray
+from .._utils import ndarray
 
 
 class VectorAdapter:

pgvector/sparsevec.py

@@ -1,12 +1,11 @@
 from __future__ import annotations
 from struct import pack, unpack_from
 from typing import TYPE_CHECKING, Any, overload
-from ._utils import is_sparse_array
+from ._utils import is_sparse_array, ndarray
 
 if TYPE_CHECKING:
     import numpy as np
     from scipy.sparse import sparray, spmatrix, coo_array, coo_matrix
-    from ._utils import ndarray
 
 
 NO_DEFAULT = object()

pgvector/vector.py

@@ -3,11 +3,10 @@
 import struct
 import sys
 from typing import TYPE_CHECKING
-from ._utils import is_ndarray
+from ._utils import is_ndarray, ndarray
 
 if TYPE_CHECKING:
     import numpy as np
-    from ._utils import ndarray
 
 
 class Vector: