stubs for the compiled module, and a gate that keeps them true

.github/workflows/ci.yml

@@ -46,5 +46,8 @@ jobs:
       # gets: the extension out of a wheel, the package out of
       # site-packages, and nothing resolved out of the checkout.
       - run: pip install ".[all]"
-      - run: pip install pytest
+      # griffe reads the stub and inspects the installed extension, so
+      # the check runs against the wheel rather than against the
+      # checkout it was built from.
+      - run: pip install pytest griffe
       - run: pytest

README.md

@@ -60,17 +60,23 @@ A result is rows to iterate and columns to hand to something else. The columns g
 
 ```python
 result = conn.execute("MATCH (p:person) RETURN p.name AS name, p.score AS score")
-result.to_arrow()      # pyarrow.Table
-result.to_pandas()     # DataFrame with Arrow-backed dtypes
-result.to_polars()     # polars.DataFrame
+result.to_arrow()  # pyarrow.Table
+result.to_pandas()  # DataFrame with Arrow-backed dtypes
+result.to_polars()  # polars.DataFrame
 result.record_batches()  # a reader, for a result larger than memory
 ```
 
 `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.
 
+## Types
+
+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.
+
+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.
+
 ## What works today
 
-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.
+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.
 
 ## Wheels
 

pyproject.toml

@@ -44,6 +44,9 @@ dev = [
   "maturin>=1.14,<2.0",
   "pytest>=8",
   "ruff>=0.9",
+  # Reads the stub as text and the built extension by inspection, which
+  # is what checks one against the other.
+  "griffe>=2",
   "pyarrow>=14",
   "pandas>=2.0",
   "polars>=1.3",

python/zudb/__init__.py

@@ -35,6 +35,7 @@
     SyntaxError,
     TransactionError,
 )
+from .types import Value
 
 __version__ = "0.0.1"
 
@@ -47,6 +48,7 @@
     "Rel",
     "Path",
     "Duration",
+    "Value",
     "Error",
     "ConnectionError",
     "DataError",

python/zudb/_zudb.pyi

@@ -0,0 +1,196 @@
+"""Types for the compiled module.
+
+The extension is a shared object, so nothing can read a signature out of
+it: mypy, pyright and an editor's completion all read this file
+instead. It ships inside the wheel next to `py.typed`, and CI checks it
+against the module it describes with griffe, which loads this file
+statically and the compiled module by inspection and compares them. A
+stub that promises a method the engine does not have is a lie a checker
+would believe, so it is worth a gate.
+
+The docstrings here are the first line of each one in the Rust source,
+because a stub is what an editor shows and an editor showing nothing is
+what a stub is for.
+"""
+
+from __future__ import annotations
+
+import datetime
+import os
+import pathlib
+from collections.abc import Iterable, Iterator, Mapping, Sequence
+from typing import Any
+
+from .types import Value
+
+#: The revision of the C ABI this client answers to.
+__abi_version__: str
+#: The version of the engine compiled into the wheel.
+__engine_version__: str
+
+def connect(
+    path: str | os.PathLike[str],
+    *,
+    read_only: bool = False,
+    memory_limit: int | None = None,
+    threads: int | None = None,
+) -> Connection:
+    """Opens the database at `path` and connects to it."""
+
+def load(
+    path: str | os.PathLike[str],
+    *,
+    nodes: str,
+    rels: str = "rel",
+    columns: Mapping[str, Iterable[Value]] | None = None,
+    edges: Iterable[Sequence[int]] | None = None,
+    rows: int | None = None,
+) -> dict[str, int]:
+    """Writes a new database at `path` and answers what went into it."""
+
+class Connection:
+    """One connection to one database."""
+
+    @property
+    def path(self) -> pathlib.Path:
+        """The file this connection was opened on."""
+
+    @property
+    def read_only(self) -> bool:
+        """Whether it was opened read-only."""
+
+    @property
+    def closed(self) -> bool:
+        """Whether this connection is still open."""
+
+    def execute(self, statement: str, params: Mapping[str, Value] | None = None) -> Result:
+        """Runs one statement and gives back its rows."""
+
+    def sql(self, statement: str, params: Mapping[str, Value] | None = None) -> Result:
+        """The same call, named for the way it reads in a notebook."""
+
+    def close(self) -> None:
+        """Closes the connection and frees what it held."""
+
+    def __enter__(self) -> Connection: ...
+    def __exit__(self, *_exception: object) -> bool: ...
+    def __repr__(self) -> str: ...
+
+class Result:
+    """The rows a statement gave back."""
+
+    @property
+    def columns(self) -> list[str]:
+        """The column names, in the order the statement projected them."""
+
+    @property
+    def notices(self) -> list[dict[str, str]]:
+        """The warnings the statement raised, if it raised any."""
+
+    def fetchall(self) -> list[tuple[Value, ...]]:
+        """Every row, as a list of tuples."""
+
+    def fetchone(self) -> tuple[Value, ...] | None:
+        """The next row, or `None` when there are no more."""
+
+    # `Any` and not `pyarrow.Table`, because the wheel does not depend
+    # on pyarrow and a stub that imported it would fail to resolve for
+    # every caller who does not have it either.
+    def to_arrow(self) -> Any:
+        """The rows as a `pyarrow.Table`."""
+
+    def to_pandas(self) -> Any:
+        """The rows as a `pandas.DataFrame`, with Arrow-backed dtypes."""
+
+    def to_polars(self) -> Any:
+        """The rows as a `polars.DataFrame`."""
+
+    def record_batches(self) -> Any:
+        """The rows as a `pyarrow.RecordBatchReader`, a batch at a time."""
+
+    def __arrow_c_stream__(self, requested_schema: object | None = None) -> Any:
+        """The rows as an Arrow stream, for anything that speaks Arrow."""
+
+    def __len__(self) -> int: ...
+    def __iter__(self) -> Iterator[tuple[Value, ...]]: ...
+    def __repr__(self) -> str: ...
+
+class Node:
+    """One node of the graph."""
+
+    def __init__(self, table: str, offset: int) -> None: ...
+    @property
+    def table(self) -> str:
+        """The name the table was given in the schema."""
+
+    @property
+    def offset(self) -> int:
+        """The row this node sits at in that table, counting from zero."""
+
+    def __hash__(self) -> int: ...
+    def __repr__(self) -> str: ...
+
+class Rel:
+    """One edge of the graph."""
+
+    def __init__(self, table: str, src: int, dst: int, ord: int) -> None: ...
+    @property
+    def table(self) -> str:
+        """The name the rel table was given in the schema."""
+
+    @property
+    def src(self) -> int:
+        """The row the edge leaves, in the node table it joins."""
+
+    @property
+    def dst(self) -> int:
+        """The row the edge arrives at."""
+
+    @property
+    def ord(self) -> int:
+        """Where the edge's properties sit, which is its place in
+        the order the table was loaded in."""
+
+    def __hash__(self) -> int: ...
+    def __repr__(self) -> str: ...
+
+class Path:
+    """A walk: nodes and edges alternating, a node at each end."""
+
+    def __init__(self, elements: list[Node | Rel]) -> None: ...
+    @property
+    def elements(self) -> list[Node | Rel]:
+        """The walk as it is stored, a node and an edge at a time."""
+
+    @property
+    def nodes(self) -> list[Node]:
+        """The nodes of the walk, in the order it visits them."""
+
+    @property
+    def rels(self) -> list[Rel]:
+        """The edges of the walk, in the order it crosses them."""
+
+    def __len__(self) -> int: ...
+    def __repr__(self) -> str: ...
+
+class Duration:
+    """A duration, which Python has no type for."""
+
+    def __init__(self, months: int = 0, nanoseconds: int = 0) -> None: ...
+    @property
+    def months(self) -> int:
+        """Months, for a year-month duration. Zero for a day-time one."""
+
+    @property
+    def nanoseconds(self) -> int:
+        """Nanoseconds, for a day-time duration. Zero for a year-month one."""
+
+    @property
+    def kind(self) -> str:
+        """`"year_month"` or `"day_time"`."""
+
+    def to_timedelta(self) -> datetime.timedelta:
+        """The same duration as a `datetime.timedelta`, rounded towards zero."""
+
+    def __hash__(self) -> int: ...
+    def __repr__(self) -> str: ...

python/zudb/types.py

@@ -0,0 +1,40 @@
+"""The type of a value, for a caller who annotates.
+
+One name for the union a row holds and a parameter takes, so a function
+that passes rows around can say what it passes. Written here rather than
+in the stub for the compiled module because a type alias a checker knows
+and the interpreter does not is a name that fails at the first
+``from zudb import Value``.
+"""
+
+from __future__ import annotations
+
+import datetime
+from typing import TypeAlias
+
+from ._zudb import Duration, Node, Path, Rel
+
+__all__ = ["Value"]
+
+#: What a statement gives back and what a parameter may be. Recursive,
+#: because a list holds values and one of them may be a list. A
+#: ``timedelta`` goes in and never comes out: zu stores it as a day-time
+#: duration and hands one back, since a ``timedelta`` cannot hold every
+#: duration zu can.
+Value: TypeAlias = (
+    None
+    | bool
+    | int
+    | float
+    | str
+    | datetime.date
+    | datetime.time
+    | datetime.datetime
+    | datetime.timedelta
+    | Duration
+    | Node
+    | Rel
+    | Path
+    | list["Value"]
+    | dict[str, "Value"]
+)