Skip to content

Commit 7de168a

Browse files
authored
API docs: fixes around vector types (#1247)
* Fix docs not resolving typing-extensions and other 3rd party lib references * Improve vector docs cross-references
1 parent a948ef5 commit 7de168a

File tree

2 files changed

+53
-40
lines changed

2 files changed

+53
-40
lines changed

pyproject.toml

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -113,8 +113,8 @@ test = [
113113
# type checker and type stubs for static type checking
114114
typing = [
115115
"mypy >= 1.15.0",
116-
"typing-extensions >= 4.13.2",
117116
"types-pytz >= 2025.2.0.20250326",
117+
{include-group = "dep-typing-extensions"},
118118
{include-group = "dep-project-dependencies"},
119119
# tests are also type-checked
120120
{include-group = "test"},
@@ -124,7 +124,11 @@ typing = [
124124
{include-group = "benchkit"},
125125
]
126126
# generate documentation
127-
docs = ["Sphinx >= 8.1.3"]
127+
docs = [
128+
"Sphinx >= 8.1.3",
129+
{include-group = "dep-typing-extensions"},
130+
{include-group = "dep-project-dependencies"},
131+
]
128132
# running the testkit backend
129133
testkit = [
130134
{include-group = "dep-freezegun"},
@@ -141,6 +145,7 @@ release = [
141145
]
142146

143147
# single dependencies and other include-groups (not really meant to be installed as a group, but to avoid duplication)
148+
dep-typing-extensions = ["typing-extensions >= 4.13.2"]
144149
dep-freezegun = ["freezegun >= 1.5.1"]
145150
dep-project-dependencies = [
146151
"pytz",

src/neo4j/vector.py

Lines changed: 46 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,18 @@
2323
from enum import Enum as _Enum
2424

2525
from . import _typing as _t
26-
from ._optional_deps import (
27-
np as _np,
28-
pa as _pa,
29-
)
26+
27+
28+
if _t.TYPE_CHECKING:
29+
# "Why?", I hear you ask. Because sphinx of course.
30+
# This beautiful construct helps sphinx to properly resolve the type hints.
31+
import numpy as _np
32+
import pyarrow as _pa
33+
else:
34+
from ._optional_deps import (
35+
np as _np,
36+
pa as _pa,
37+
)
3038

3139

3240
if False:
@@ -41,10 +49,6 @@
4149
_swap_endian_unchecked_rust = None
4250
_vec_rust = None
4351

44-
if _t.TYPE_CHECKING:
45-
import numpy # type: ignore[import]
46-
import pyarrow # type: ignore[import]
47-
4852

4953
__all__ = [
5054
"Vector",
@@ -84,14 +88,18 @@ class Vector:
8488
Use an iterable of floats or an iterable of ints to construct the
8589
vector from native Python values.
8690
The ``dtype`` parameter is required.
91+
See also: :meth:`.from_native`.
8792
* ``bytes``, ``bytearray``: Use raw bytes to construct the vector.
8893
The ``dtype`` parameter is required and ``byteorder`` is optional.
8994
* ``numpy.ndarray``: Use a numpy array to construct the vector.
9095
No further parameters are accepted.
96+
See also: :meth:`.from_numpy`.
9197
* ``pyarrow.Array``: Use a pyarrow array to construct the vector.
9298
No further parameters are accepted.
99+
See also: :meth:`.from_pyarrow`.
93100
:param dtype: The type of the vector.
94-
See :attr:`.dtype` for currently supported inner data types.
101+
See :class:`.VectorDType` for currently supported inner data types.
102+
See also :attr:`.dtype`.
95103
96104
This parameter is required if ``data`` is of type :class:`bytes`,
97105
:class:`bytearray`, ``Iterable[float]``, or ``Iterable[int]``.
@@ -163,10 +171,10 @@ def __init__(
163171
) -> None: ...
164172

165173
@_t.overload
166-
def __init__(self, data: numpy.ndarray, /) -> None: ...
174+
def __init__(self, data: _np.ndarray, /) -> None: ...
167175

168176
@_t.overload
169-
def __init__(self, data: pyarrow.Array, /) -> None: ...
177+
def __init__(self, data: _pa.Array, /) -> None: ...
170178

171179
def __init__(self, data, *args, **kwargs) -> None:
172180
if isinstance(data, (bytes, bytearray)):
@@ -373,15 +381,15 @@ def to_native(self) -> list[object]:
373381
return self._inner.to_native()
374382

375383
@classmethod
376-
def from_numpy(cls, data: numpy.ndarray, /) -> _t.Self:
384+
def from_numpy(cls, data: _np.ndarray, /) -> _t.Self:
377385
"""
378386
Create a Vector instance from a numpy array.
379387
380388
:param data: The numpy array to create the vector from.
381389
The array must be one-dimensional and have a dtype that is
382390
supported by Neo4j vectors: ``float64``, ``float32``,
383391
``int64``, ``int32``, ``int16``, or ``int8``.
384-
See also :attr:`.dtype`.
392+
See also :class:`.VectorDType`.
385393
386394
:raises ValueError:
387395
* If the dtype is not supported.
@@ -394,7 +402,7 @@ def from_numpy(cls, data: numpy.ndarray, /) -> _t.Self:
394402
obj._set_numpy(data)
395403
return obj
396404

397-
def to_numpy(self) -> numpy.ndarray:
405+
def to_numpy(self) -> _np.ndarray:
398406
"""
399407
Convert the vector to a numpy array.
400408
@@ -407,7 +415,7 @@ def to_numpy(self) -> numpy.ndarray:
407415
"""
408416
return self._inner.to_numpy()
409417

410-
def _set_numpy(self, data: numpy.ndarray, /) -> None:
418+
def _set_numpy(self, data: _np.ndarray, /) -> None:
411419
if data.ndim != 1:
412420
raise ValueError("Data must be one-dimensional")
413421
type_: type[_InnerVector]
@@ -429,18 +437,17 @@ def _set_numpy(self, data: numpy.ndarray, /) -> None:
429437
self._inner = type_.from_numpy(data)
430438

431439
@classmethod
432-
def from_pyarrow(cls, data: pyarrow.Array, /) -> _t.Self:
440+
def from_pyarrow(cls, data: _pa.Array, /) -> _t.Self:
433441
"""
434442
Create a Vector instance from a pyarrow array.
435443
436-
:param data: The pyarrow array to create the vector from.
437-
The array must have a type that is supported by Neo4j.
438-
See also :attr:`.dtype`.
439-
440444
PyArrow stores data in little endian. Therefore, the byte-order needs
441445
to be swapped. If ``neo4j-rust-ext`` or ``numpy`` is installed, it will
442446
be used to speed up the byte flipping.
443447
448+
:param data: The pyarrow array to create the vector from.
449+
The array must have a type that is supported by Neo4j.
450+
See also :class:`.VectorDType`.
444451
:raises ValueError:
445452
* If the array's type is not supported.
446453
* If the array contains null values.
@@ -452,7 +459,7 @@ def from_pyarrow(cls, data: pyarrow.Array, /) -> _t.Self:
452459
obj._set_pyarrow(data)
453460
return obj
454461

455-
def to_pyarrow(self) -> pyarrow.Array:
462+
def to_pyarrow(self) -> _pa.Array:
456463
"""
457464
Convert the vector to a pyarrow array.
458465
@@ -462,7 +469,7 @@ def to_pyarrow(self) -> pyarrow.Array:
462469
"""
463470
return self._inner.to_pyarrow()
464471

465-
def _set_pyarrow(self, data: pyarrow.Array, /) -> None:
472+
def _set_pyarrow(self, data: _pa.Array, /) -> None:
466473
import pyarrow
467474

468475
type_: type[_InnerVector]
@@ -581,6 +588,7 @@ def _swap_endian(type_size: int, data: bytes, /) -> bytes:
581588

582589

583590
def _swap_endian_unchecked_np(type_size: int, data: bytes, /) -> bytes:
591+
dtype: _np.dtype
584592
match type_size:
585593
case 2:
586594
dtype = _np.dtype("<i2")
@@ -727,18 +735,18 @@ def from_native(cls, data: _t.Iterable[object], /) -> _t.Self: ...
727735
def to_native(self) -> list[object]: ...
728736

729737
@classmethod
730-
def from_numpy(cls, data: numpy.ndarray, /) -> _t.Self:
738+
def from_numpy(cls, data: _np.ndarray, /) -> _t.Self:
731739
if data.dtype.byteorder == "<" or (
732740
data.dtype.byteorder == "=" and _sys.byteorder == "little"
733741
):
734742
data = data.byteswap()
735743
return cls(data.tobytes())
736744

737745
@_abc.abstractmethod
738-
def to_numpy(self) -> numpy.ndarray: ...
746+
def to_numpy(self) -> _np.ndarray: ...
739747

740748
@classmethod
741-
def from_pyarrow(cls, data: pyarrow.Array, /) -> _t.Self:
749+
def from_pyarrow(cls, data: _pa.Array, /) -> _t.Self:
742750
width = data.type.byte_width
743751
assert cls.size == width
744752
if _pa.compute.count(data, mode="only_null").as_py():
@@ -750,7 +758,7 @@ def from_pyarrow(cls, data: pyarrow.Array, /) -> _t.Self:
750758
return cls(bytes(buffer), byteorder=_sys.byteorder)
751759

752760
@_abc.abstractmethod
753-
def to_pyarrow(self) -> pyarrow.Array: ...
761+
def to_pyarrow(self) -> _pa.Array: ...
754762

755763

756764
class _InnerVectorFloat(_InnerVector, _abc.ABC):
@@ -822,12 +830,12 @@ def _to_native_py(self) -> list[object]:
822830
else:
823831
to_native = _to_native_py
824832

825-
def to_numpy(self) -> numpy.ndarray:
833+
def to_numpy(self) -> _np.ndarray:
826834
import numpy
827835

828836
return numpy.frombuffer(self.data, dtype=numpy.dtype(">f8"))
829837

830-
def to_pyarrow(self) -> pyarrow.Array:
838+
def to_pyarrow(self) -> _pa.Array:
831839
import pyarrow
832840

833841
buffer = pyarrow.py_buffer(self.data_le)
@@ -897,12 +905,12 @@ def _to_native_py(self) -> list[object]:
897905
else:
898906
to_native = _to_native_py
899907

900-
def to_numpy(self) -> numpy.ndarray:
908+
def to_numpy(self) -> _np.ndarray:
901909
import numpy
902910

903911
return numpy.frombuffer(self.data, dtype=numpy.dtype(">f4"))
904912

905-
def to_pyarrow(self) -> pyarrow.Array:
913+
def to_pyarrow(self) -> _pa.Array:
906914
import pyarrow
907915

908916
buffer = pyarrow.py_buffer(self.data_le)
@@ -997,12 +1005,12 @@ def _to_native_py(self) -> list[object]:
9971005
else:
9981006
to_native = _to_native_py
9991007

1000-
def to_numpy(self) -> numpy.ndarray:
1008+
def to_numpy(self) -> _np.ndarray:
10011009
import numpy
10021010

10031011
return numpy.frombuffer(self.data, dtype=numpy.dtype(">i8"))
10041012

1005-
def to_pyarrow(self) -> pyarrow.Array:
1013+
def to_pyarrow(self) -> _pa.Array:
10061014
import pyarrow
10071015

10081016
buffer = pyarrow.py_buffer(self.data_le)
@@ -1090,12 +1098,12 @@ def _to_native_py(self) -> list[object]:
10901098
else:
10911099
to_native = _to_native_py
10921100

1093-
def to_numpy(self) -> numpy.ndarray:
1101+
def to_numpy(self) -> _np.ndarray:
10941102
import numpy
10951103

10961104
return numpy.frombuffer(self.data, dtype=numpy.dtype(">i4"))
10971105

1098-
def to_pyarrow(self) -> pyarrow.Array:
1106+
def to_pyarrow(self) -> _pa.Array:
10991107
import pyarrow
11001108

11011109
buffer = pyarrow.py_buffer(self.data_le)
@@ -1183,12 +1191,12 @@ def _to_native_py(self) -> list[object]:
11831191
else:
11841192
to_native = _to_native_py
11851193

1186-
def to_numpy(self) -> numpy.ndarray:
1194+
def to_numpy(self) -> _np.ndarray:
11871195
import numpy
11881196

11891197
return numpy.frombuffer(self.data, dtype=numpy.dtype(">i2"))
11901198

1191-
def to_pyarrow(self) -> pyarrow.Array:
1199+
def to_pyarrow(self) -> _pa.Array:
11921200
import pyarrow
11931201

11941202
buffer = pyarrow.py_buffer(self.data_le)
@@ -1276,12 +1284,12 @@ def _to_native_py(self) -> list[object]:
12761284
else:
12771285
to_native = _to_native_py
12781286

1279-
def to_numpy(self) -> numpy.ndarray:
1287+
def to_numpy(self) -> _np.ndarray:
12801288
import numpy
12811289

12821290
return numpy.frombuffer(self.data, dtype=numpy.dtype(">i1"))
12831291

1284-
def to_pyarrow(self) -> pyarrow.Array:
1292+
def to_pyarrow(self) -> _pa.Array:
12851293
import pyarrow
12861294

12871295
buffer = pyarrow.py_buffer(self.data_le)

0 commit comments

Comments
 (0)