Skip to content

Commit bc20cc0

Browse files
authored
stubs for the compiled module, and a gate that keeps them true (#4)
An extension module is a shared object and nothing can read a signature out of one, so mypy, pyright and an editor's completion all had nothing to go on. This is `_zudb.pyi`: every function, class, property and method, with the parameter names and defaults the module actually has. It ships in the wheel beside `py.typed`, which was already there. `zudb.Value` comes with it, the union a row holds and a parameter takes, for code that passes rows around and wants to say so. Written as a real module rather than only in the stub, because a type alias a checker knows and the interpreter does not is a name that fails at the first `from zudb import Value`. A stub is a promise no interpreter checks. Nothing fails when it names a method the engine does not have or gives a parameter a name the engine does not answer to; what fails is the caller who believed the editor. So it is checked: griffe reads the stub as text, inspects the installed extension, and the two are compared on every public name, every kind, every parameter and every default. PyO3 writes a text signature for everything it exports, which is what makes the second half possible. Two things the comparison has to work around. griffe merges a stub into the extension beside it, so a merged stub would agree with the module by construction and a missing name would come back from the `.so`; the package is copied to a temporary directory without its binary first. And a constructor is exported as `__new__` with `(*args, **kwargs)`, with the real signature on the class, where `inspect` finds it and where the stub's `__init__` is checked against it. What is not compared is the types. No inspection of a compiled module can confirm that a column is a `list[str]`, and the tests that call it are what say so. 25 tests for the stub, 176 in the client. Verified by breaking it both ways: a method the module does not have and a parameter renamed, each caught by the test that should catch it.
1 parent b8ce46c commit bc20cc0

7 files changed

Lines changed: 461 additions & 5 deletions

File tree

‎.github/workflows/ci.yml‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,5 +46,8 @@ jobs:
4646
# gets: the extension out of a wheel, the package out of
4747
# site-packages, and nothing resolved out of the checkout.
4848
- run: pip install ".[all]"
49-
- run: pip install pytest
49+
# griffe reads the stub and inspects the installed extension, so
50+
# the check runs against the wheel rather than against the
51+
# checkout it was built from.
52+
- run: pip install pytest griffe
5053
- run: pytest

‎README.md‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -60,17 +60,23 @@ A result is rows to iterate and columns to hand to something else. The columns g
6060

6161
```python
6262
result = conn.execute("MATCH (p:person) RETURN p.name AS name, p.score AS score")
63-
result.to_arrow() # pyarrow.Table
64-
result.to_pandas() # DataFrame with Arrow-backed dtypes
65-
result.to_polars() # polars.DataFrame
63+
result.to_arrow() # pyarrow.Table
64+
result.to_pandas() # DataFrame with Arrow-backed dtypes
65+
result.to_polars() # polars.DataFrame
6666
result.record_batches() # a reader, for a result larger than memory
6767
```
6868

6969
`Result` implements `__arrow_c_stream__`, so anything that reads the protocol reads a result directly and none of the four methods above is needed: `pyarrow.table(result)` and `polars.DataFrame(result)` both work. Batches are 65,536 rows. A column holds one type, which the values decide, and integers beside floats are the one mixture that widens rather than being refused. Nodes, rels and paths go across as structs. The copy runs with the GIL released, and on this machine 300,000 rows across three columns take 44 ms as Arrow against 67 ms as Python objects, and a single integer column takes 13.8 ms against 44.5 ms.
7070

71+
## Types
72+
73+
The wheel carries `py.typed` and a stub for the compiled module, so mypy, pyright and an editor's completion all work with nothing else installed. `zudb.Value` is the union a row holds and a parameter takes, for code that passes rows around and wants to say so.
74+
75+
The stub is checked against the module it describes in CI: griffe reads the stub as text and the installed extension by inspection, and the two have to agree on every name, every parameter and every default. A stub is a promise no interpreter checks, so something has to.
76+
7177
## What works today
7278

73-
The list above is what this client is for. What it does so far is the core of it: `connect`, `execute` and `sql` with named parameters, results that iterate and fetch, values as Python objects both ways including dates, times, datetimes and durations, `Node`, `Rel` and `Path` as classes, `load` for building a graph with edges in it, every condition as an exception class carrying its code, its position and its documentation link, results as Arrow columns and as pandas and polars frames, and the GIL released around every statement, every load and every copy out. `register`, the stubs and the interrupt are next, and each one lands with the tests that say it works.
79+
The list above is what this client is for. What it does so far is the core of it: `connect`, `execute` and `sql` with named parameters, results that iterate and fetch, values as Python objects both ways including dates, times, datetimes and durations, `Node`, `Rel` and `Path` as classes, `load` for building a graph with edges in it, every condition as an exception class carrying its code, its position and its documentation link, results as Arrow columns and as pandas and polars frames, stubs inside the wheel with a gate that keeps them true, and the GIL released around every statement, every load and every copy out. `register` and the interrupt are next, and each one lands with the tests that say it works.
7480

7581
## Wheels
7682

‎pyproject.toml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,9 @@ dev = [
4444
"maturin>=1.14,<2.0",
4545
"pytest>=8",
4646
"ruff>=0.9",
47+
# Reads the stub as text and the built extension by inspection, which
48+
# is what checks one against the other.
49+
"griffe>=2",
4750
"pyarrow>=14",
4851
"pandas>=2.0",
4952
"polars>=1.3",

‎python/zudb/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
SyntaxError,
3636
TransactionError,
3737
)
38+
from .types import Value
3839

3940
__version__ = "0.0.1"
4041

@@ -47,6 +48,7 @@
4748
"Rel",
4849
"Path",
4950
"Duration",
51+
"Value",
5052
"Error",
5153
"ConnectionError",
5254
"DataError",

‎python/zudb/_zudb.pyi‎

Lines changed: 196 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,196 @@
1+
"""Types for the compiled module.
2+
3+
The extension is a shared object, so nothing can read a signature out of
4+
it: mypy, pyright and an editor's completion all read this file
5+
instead. It ships inside the wheel next to `py.typed`, and CI checks it
6+
against the module it describes with griffe, which loads this file
7+
statically and the compiled module by inspection and compares them. A
8+
stub that promises a method the engine does not have is a lie a checker
9+
would believe, so it is worth a gate.
10+
11+
The docstrings here are the first line of each one in the Rust source,
12+
because a stub is what an editor shows and an editor showing nothing is
13+
what a stub is for.
14+
"""
15+
16+
from __future__ import annotations
17+
18+
import datetime
19+
import os
20+
import pathlib
21+
from collections.abc import Iterable, Iterator, Mapping, Sequence
22+
from typing import Any
23+
24+
from .types import Value
25+
26+
#: The revision of the C ABI this client answers to.
27+
__abi_version__: str
28+
#: The version of the engine compiled into the wheel.
29+
__engine_version__: str
30+
31+
def connect(
32+
path: str | os.PathLike[str],
33+
*,
34+
read_only: bool = False,
35+
memory_limit: int | None = None,
36+
threads: int | None = None,
37+
) -> Connection:
38+
"""Opens the database at `path` and connects to it."""
39+
40+
def load(
41+
path: str | os.PathLike[str],
42+
*,
43+
nodes: str,
44+
rels: str = "rel",
45+
columns: Mapping[str, Iterable[Value]] | None = None,
46+
edges: Iterable[Sequence[int]] | None = None,
47+
rows: int | None = None,
48+
) -> dict[str, int]:
49+
"""Writes a new database at `path` and answers what went into it."""
50+
51+
class Connection:
52+
"""One connection to one database."""
53+
54+
@property
55+
def path(self) -> pathlib.Path:
56+
"""The file this connection was opened on."""
57+
58+
@property
59+
def read_only(self) -> bool:
60+
"""Whether it was opened read-only."""
61+
62+
@property
63+
def closed(self) -> bool:
64+
"""Whether this connection is still open."""
65+
66+
def execute(self, statement: str, params: Mapping[str, Value] | None = None) -> Result:
67+
"""Runs one statement and gives back its rows."""
68+
69+
def sql(self, statement: str, params: Mapping[str, Value] | None = None) -> Result:
70+
"""The same call, named for the way it reads in a notebook."""
71+
72+
def close(self) -> None:
73+
"""Closes the connection and frees what it held."""
74+
75+
def __enter__(self) -> Connection: ...
76+
def __exit__(self, *_exception: object) -> bool: ...
77+
def __repr__(self) -> str: ...
78+
79+
class Result:
80+
"""The rows a statement gave back."""
81+
82+
@property
83+
def columns(self) -> list[str]:
84+
"""The column names, in the order the statement projected them."""
85+
86+
@property
87+
def notices(self) -> list[dict[str, str]]:
88+
"""The warnings the statement raised, if it raised any."""
89+
90+
def fetchall(self) -> list[tuple[Value, ...]]:
91+
"""Every row, as a list of tuples."""
92+
93+
def fetchone(self) -> tuple[Value, ...] | None:
94+
"""The next row, or `None` when there are no more."""
95+
96+
# `Any` and not `pyarrow.Table`, because the wheel does not depend
97+
# on pyarrow and a stub that imported it would fail to resolve for
98+
# every caller who does not have it either.
99+
def to_arrow(self) -> Any:
100+
"""The rows as a `pyarrow.Table`."""
101+
102+
def to_pandas(self) -> Any:
103+
"""The rows as a `pandas.DataFrame`, with Arrow-backed dtypes."""
104+
105+
def to_polars(self) -> Any:
106+
"""The rows as a `polars.DataFrame`."""
107+
108+
def record_batches(self) -> Any:
109+
"""The rows as a `pyarrow.RecordBatchReader`, a batch at a time."""
110+
111+
def __arrow_c_stream__(self, requested_schema: object | None = None) -> Any:
112+
"""The rows as an Arrow stream, for anything that speaks Arrow."""
113+
114+
def __len__(self) -> int: ...
115+
def __iter__(self) -> Iterator[tuple[Value, ...]]: ...
116+
def __repr__(self) -> str: ...
117+
118+
class Node:
119+
"""One node of the graph."""
120+
121+
def __init__(self, table: str, offset: int) -> None: ...
122+
@property
123+
def table(self) -> str:
124+
"""The name the table was given in the schema."""
125+
126+
@property
127+
def offset(self) -> int:
128+
"""The row this node sits at in that table, counting from zero."""
129+
130+
def __hash__(self) -> int: ...
131+
def __repr__(self) -> str: ...
132+
133+
class Rel:
134+
"""One edge of the graph."""
135+
136+
def __init__(self, table: str, src: int, dst: int, ord: int) -> None: ...
137+
@property
138+
def table(self) -> str:
139+
"""The name the rel table was given in the schema."""
140+
141+
@property
142+
def src(self) -> int:
143+
"""The row the edge leaves, in the node table it joins."""
144+
145+
@property
146+
def dst(self) -> int:
147+
"""The row the edge arrives at."""
148+
149+
@property
150+
def ord(self) -> int:
151+
"""Where the edge's properties sit, which is its place in
152+
the order the table was loaded in."""
153+
154+
def __hash__(self) -> int: ...
155+
def __repr__(self) -> str: ...
156+
157+
class Path:
158+
"""A walk: nodes and edges alternating, a node at each end."""
159+
160+
def __init__(self, elements: list[Node | Rel]) -> None: ...
161+
@property
162+
def elements(self) -> list[Node | Rel]:
163+
"""The walk as it is stored, a node and an edge at a time."""
164+
165+
@property
166+
def nodes(self) -> list[Node]:
167+
"""The nodes of the walk, in the order it visits them."""
168+
169+
@property
170+
def rels(self) -> list[Rel]:
171+
"""The edges of the walk, in the order it crosses them."""
172+
173+
def __len__(self) -> int: ...
174+
def __repr__(self) -> str: ...
175+
176+
class Duration:
177+
"""A duration, which Python has no type for."""
178+
179+
def __init__(self, months: int = 0, nanoseconds: int = 0) -> None: ...
180+
@property
181+
def months(self) -> int:
182+
"""Months, for a year-month duration. Zero for a day-time one."""
183+
184+
@property
185+
def nanoseconds(self) -> int:
186+
"""Nanoseconds, for a day-time duration. Zero for a year-month one."""
187+
188+
@property
189+
def kind(self) -> str:
190+
"""`"year_month"` or `"day_time"`."""
191+
192+
def to_timedelta(self) -> datetime.timedelta:
193+
"""The same duration as a `datetime.timedelta`, rounded towards zero."""
194+
195+
def __hash__(self) -> int: ...
196+
def __repr__(self) -> str: ...

‎python/zudb/types.py‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
"""The type of a value, for a caller who annotates.
2+
3+
One name for the union a row holds and a parameter takes, so a function
4+
that passes rows around can say what it passes. Written here rather than
5+
in the stub for the compiled module because a type alias a checker knows
6+
and the interpreter does not is a name that fails at the first
7+
``from zudb import Value``.
8+
"""
9+
10+
from __future__ import annotations
11+
12+
import datetime
13+
from typing import TypeAlias
14+
15+
from ._zudb import Duration, Node, Path, Rel
16+
17+
__all__ = ["Value"]
18+
19+
#: What a statement gives back and what a parameter may be. Recursive,
20+
#: because a list holds values and one of them may be a list. A
21+
#: ``timedelta`` goes in and never comes out: zu stores it as a day-time
22+
#: duration and hands one back, since a ``timedelta`` cannot hold every
23+
#: duration zu can.
24+
Value: TypeAlias = (
25+
None
26+
| bool
27+
| int
28+
| float
29+
| str
30+
| datetime.date
31+
| datetime.time
32+
| datetime.datetime
33+
| datetime.timedelta
34+
| Duration
35+
| Node
36+
| Rel
37+
| Path
38+
| list["Value"]
39+
| dict[str, "Value"]
40+
)

0 commit comments

Comments
 (0)