From 6175122bd24cc8fe95ddd4cf82d7492975ea6e7d Mon Sep 17 00:00:00 2001 From: David Hughes Date: Sun, 25 Oct 2026 09:56:11 -0700 Subject: [PATCH] feat(python): ship public API type stubs Add PEP 561 stubs matching the extension graph module, strict consumer checks, and runtime signature parity tests. Document the typing workflow and check wheel inclusion. Runtime and benchmark validation remain pending the coordinated build. --- ci/check_wheel.py | 3 + docs/api-reference.md | 29 ++++ python/DEVELOPMENT.md | 24 ++- python/pyproject.toml | 13 +- python/python/lance_csr/__init__.py | 2 +- .../python/lance_csr/_internal/__init__.pyi | 3 + python/python/lance_csr/_internal/graph.pyi | 155 ++++++++++++++++++ python/python/lance_csr/py.typed | 0 python/python/tests/test_typing.py | 149 +++++++++++++++++ python/python/tests/typing/public_api.py | 95 +++++++++++ python/uv.lock | 14 ++ 11 files changed, 482 insertions(+), 5 deletions(-) create mode 100644 python/python/lance_csr/_internal/__init__.pyi create mode 100644 python/python/lance_csr/_internal/graph.pyi create mode 100644 python/python/lance_csr/py.typed create mode 100644 python/python/tests/test_typing.py create mode 100644 python/python/tests/typing/public_api.py diff --git a/ci/check_wheel.py b/ci/check_wheel.py index 3d3e31c..4a5999e 100644 --- a/ci/check_wheel.py +++ b/ci/check_wheel.py @@ -36,6 +36,9 @@ def main(): # outside sys.prefix, so check for site-packages itself). extension = Path(_internal.__file__).resolve() assert "site-packages" in extension.parts, extension + package = extension.parent + for relative in ("py.typed", "_internal/__init__.pyi", "_internal/graph.pyi"): + assert (package / relative).is_file(), f"Missing typing file: {relative}" config = ( GraphConfig.builder() diff --git a/docs/api-reference.md b/docs/api-reference.md index 81ffd03..9f5289b 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -16,6 +16,35 @@ code. The guides link here for exact signatures and link back for examples. `lance_csr.__version__` is the package version. +## Type checking + +The package ships type stubs for every public class, including the compiled extension, and a `py.typed` +marker. Editors and type checkers can follow builder methods, parameters, query results and statistics. +The stubs describe the current API; they do not change runtime validation or exceptions. + +PyArrow distributes its type information separately. For strict checks of Arrow inputs and results, add +`pyarrow-stubs` to your development dependencies: + +```bash +uv add --dev pyright pyarrow-stubs +uv run pyright your_script.py +``` + +```python +from typing import assert_type + +import pyarrow as pa +from lance_csr import CypherEngine, GraphConfig + +config = GraphConfig.builder().node_table("Person", "id").build() +engine = CypherEngine(config, {"Person": pa.table({"id": [1, 2]})}) +result = engine.execute("MATCH (p:Person) RETURN p.id ORDER BY p.id") +assert_type(result, pa.Table) +print(result.to_pylist()) # [{'p.id': 1}, {'p.id': 2}] +``` + +This reference remains maintained by hand alongside the stubs and Rust bindings. + ## Datasets Methods that take `datasets` take a `dict[str, object]` mapping each node label or relationship type (or SQL diff --git a/python/DEVELOPMENT.md b/python/DEVELOPMENT.md index 885a34c..87d7401 100644 --- a/python/DEVELOPMENT.md +++ b/python/DEVELOPMENT.md @@ -20,8 +20,8 @@ uv sync --locked --no-install-project # locked dependencies and the dev group uv run --no-sync maturin develop # debug build of the extension, installed in editable mode ``` -`make install` runs both. The `dev` dependency group in `pyproject.toml` holds maturin, pytest, pandas, ruff and -pyright, and `uv sync` installs it by default. +`make install` runs both. The `dev` dependency group in `pyproject.toml` holds maturin, pytest, pandas, ruff, +pyright and PyArrow's separate stubs, and `uv sync` installs it by default. - `--no-install-project` leaves the package to `maturin develop`. Without it, uv builds the package through maturin's build backend, in release mode. @@ -59,6 +59,26 @@ make format # ruff format make lint # ruff check, ruff format --check, pyright ``` +The public Python API is typed in `python/lance_csr/_internal/graph.pyi`, matching the compiled +`_internal.graph` module. `_internal/__init__.pyi` describes its parent module, and `py.typed` marks the +package as typed. Update the stubs, the hand-written API reference and the relevant guide whenever a +binding changes. Exceptions remain the built-in Python exceptions listed in the API reference. + +`make lint` checks `python/tests/typing/public_api.py` in strict mode: it exercises every class, checks +return types, and confirms invalid calls remain rejected. PyArrow's separate stubs are a development +dependency because Arrow tables and readers are part of this API. They are not a runtime dependency. + +After building the extension, check public members, enum variants, properties, constructors, argument +names and defaults against the stubs: + +```shell +uv run --no-sync pytest python/tests/test_typing.py -q +``` + +These tests fail when an exported class or public method has no stub. They also check that each top-level +export refers to the class declared by the extension. The wheel smoke check verifies that the stubs and +marker survive packaging; run it against the installed wheel as described below. + For Rust, from the repository root (CI runs the same checks): ```shell diff --git a/python/pyproject.toml b/python/pyproject.toml index bb2c8fa..10b9e3a 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -40,7 +40,14 @@ build-backend = "maturin" # Development tools, installed by `uv sync` (the `dev` group is uv's default). [dependency-groups] -dev = ["maturin[patchelf]>=1.9,<2", "pytest", "pandas", "ruff", "pyright"] +dev = [ + "maturin[patchelf]>=1.9,<2", + "pytest", + "pandas", + "ruff", + "pyright", + "pyarrow-stubs>=20.0.0.20260819", +] [tool.ruff] lint.select = ["F", "E", "W", "I", "G", "TCH", "PERF", "B019"] @@ -50,7 +57,9 @@ lint.select = ["F", "E", "W", "I", "G", "TCH", "PERF", "B019"] [tool.pyright] pythonVersion = "3.13" -include = ["python/lance_csr/__init__.py"] +include = ["python/lance_csr", "python/tests/typing"] +strict = ["python/tests/typing"] +reportUnnecessaryTypeIgnoreComment = "error" # Dependencies like pyarrow make this difficult to enforce strictly. reportMissingTypeStubs = "warning" reportImportCycles = "error" diff --git a/python/python/lance_csr/__init__.py b/python/python/lance_csr/__init__.py index 196aa72..8a29ec4 100644 --- a/python/python/lance_csr/__init__.py +++ b/python/python/lance_csr/__init__.py @@ -11,7 +11,7 @@ from __future__ import annotations try: - from . import _internal as _bindings # type: ignore[attr-defined] + from . import _internal as _bindings except ImportError as e: raise ImportError( "Unable to import the compiled lance-csr extension. " diff --git a/python/python/lance_csr/_internal/__init__.pyi b/python/python/lance_csr/_internal/__init__.pyi new file mode 100644 index 0000000..a2ac205 --- /dev/null +++ b/python/python/lance_csr/_internal/__init__.pyi @@ -0,0 +1,3 @@ +from . import graph as graph + +__version__: str diff --git a/python/python/lance_csr/_internal/graph.pyi b/python/python/lance_csr/_internal/graph.pyi new file mode 100644 index 0000000..1c6b13a --- /dev/null +++ b/python/python/lance_csr/_internal/graph.pyi @@ -0,0 +1,155 @@ +from collections.abc import Sequence +from datetime import date, datetime +from typing import ClassVar, Literal, Protocol, Self, TypedDict, TypeVar, final + +from lance import LanceDataset +from pyarrow import RecordBatch, RecordBatchReader, Table + +class _ArrowStream(Protocol): + def __arrow_c_stream__(self, requested_schema: object = None) -> object: ... + +_Dataset = Table | RecordBatch | RecordBatchReader | LanceDataset | _ArrowStream +_DatasetT = TypeVar("_DatasetT", bound=_Dataset) +_Parameter = int | float | str | bool | date | datetime | None +_ParameterT = TypeVar("_ParameterT", bound=_Parameter) +_Scalar = int | float | str | bool | None + +class _NodeIndexStats(TypedDict): + rows: int + version: int | None + build_ms: float + +class _CsrStats(TypedDict): + edges: int + delta_edges: int + tombstones: int + update: Literal["build", "rebuild", "delta", "reused"] + version: int | None + build_ms: float + +@final +class ExecutionStrategy: + DataFusion: ClassVar[ExecutionStrategy] + LanceNative: ClassVar[ExecutionStrategy] + def __int__(self) -> int: ... + +@final +class SqlDialect: + Default: ClassVar[SqlDialect] + Spark: ClassVar[SqlDialect] + PostgreSql: ClassVar[SqlDialect] + MySql: ClassVar[SqlDialect] + Sqlite: ClassVar[SqlDialect] + def __int__(self) -> int: ... + +@final +class DistanceMetric: + L2: ClassVar[DistanceMetric] + Cosine: ClassVar[DistanceMetric] + Dot: ClassVar[DistanceMetric] + def __int__(self) -> int: ... + +@final +class GraphConfig: + @staticmethod + def builder() -> GraphConfigBuilder: ... + def node_labels(self) -> list[str]: ... + def relationship_types(self) -> list[str]: ... + +@final +class GraphConfigBuilder: + def __init__(self) -> None: ... + def node_table(self, label: str, primary_key: str) -> Self: ... + def rel_table( + self, rel_type: str, from_table: str, to_table: str, src: str, dst: str + ) -> Self: ... + def with_node_label(self, label: str, id_field: str) -> Self: ... + def with_relationship( + self, + rel_type: str, + source_field: str, + target_field: str, + source_label: str | None = None, + target_label: str | None = None, + ) -> Self: ... + def build(self) -> GraphConfig: ... + +@final +class CypherQuery: + def __init__(self, query_text: str) -> None: ... + def with_config(self, config: GraphConfig) -> Self: ... + def with_parameter(self, key: str, value: _Parameter) -> Self: ... + def query_text(self) -> str: ... + def parameters(self) -> dict[str, _Scalar]: ... + def to_sql( + self, datasets: dict[str, _DatasetT], dialect: SqlDialect | None = None + ) -> str: ... + def execute( + self, + datasets: dict[str, _DatasetT], + strategy: ExecutionStrategy | None = None, + ) -> Table: ... + def execute_with_namespace( + self, namespace: DirNamespace, strategy: ExecutionStrategy | None = None + ) -> Table: ... + def explain(self, datasets: dict[str, _DatasetT]) -> str: ... + def variables(self) -> list[str]: ... + def node_labels(self) -> list[str]: ... + def relationship_types(self) -> list[str]: ... + def execute_with_vector_rerank( + self, datasets: dict[str, _DatasetT], vector_search: VectorSearch + ) -> Table: ... + +@final +class CypherEngine: + def __init__( + self, config: GraphConfig, datasets: dict[str, _DatasetT], use_csr: bool = False + ) -> None: ... + def execute( + self, query: str, params: dict[str, _ParameterT] | None = None + ) -> Table: ... + def execute_batches( + self, query: str, params: dict[str, _ParameterT] | None = None + ) -> RecordBatchReader: ... + def explain( + self, query: str, params: dict[str, _ParameterT] | None = None + ) -> str: ... + def refresh(self) -> None: ... + def node_index_stats(self) -> dict[str, _NodeIndexStats]: ... + def csr_stats(self) -> dict[str, _CsrStats]: ... + def dataset_versions(self) -> dict[str, int]: ... + def config(self) -> GraphConfig: ... + def execute_with_vector_rerank( + self, query: str, vector_search: VectorSearch + ) -> Table: ... + +@final +class SqlQuery: + def __init__(self, sql: str) -> None: ... + def sql(self) -> str: ... + def execute(self, datasets: dict[str, _DatasetT]) -> Table: ... + def explain(self, datasets: dict[str, _DatasetT]) -> str: ... + +@final +class SqlEngine: + def __init__(self, datasets: dict[str, _DatasetT]) -> None: ... + def execute(self, sql: str) -> Table: ... + +@final +class VectorSearch: + def __init__(self, column: str) -> None: ... + def query_vector(self, vector: Sequence[float]) -> Self: ... + def metric(self, metric: DistanceMetric) -> Self: ... + def top_k(self, k: int) -> Self: ... + def include_distance(self, include: bool) -> Self: ... + def distance_column_name(self, name: str) -> Self: ... + def use_lance_index(self, enabled: bool) -> Self: ... + def search( + self, table: Table | RecordBatch | RecordBatchReader | _ArrowStream + ) -> Table: ... + +@final +class DirNamespace: + def __init__(self, base_uri: str) -> None: ... + @property + def base_uri(self) -> str: ... diff --git a/python/python/lance_csr/py.typed b/python/python/lance_csr/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/python/python/tests/test_typing.py b/python/python/tests/test_typing.py new file mode 100644 index 0000000..4087b4d --- /dev/null +++ b/python/python/tests/test_typing.py @@ -0,0 +1,149 @@ +"""Keep shipped stubs aligned with the compiled public API.""" + +import ast +import inspect +from pathlib import Path + +import lance_csr +import pytest +from lance_csr import _internal + + +def stub_classes(): + """Read the installed stubs, so the same check also exercises a built wheel.""" + package = Path(lance_csr.__file__).parent + assert (package / "py.typed").is_file() + tree = ast.parse((package / "_internal" / "graph.pyi").read_text()) + return { + node.name: node + for node in tree.body + if isinstance(node, ast.ClassDef) and not node.name.startswith("_") + } + + +def public_members(node): + return { + item.name if isinstance(item, ast.FunctionDef) else item.target.id: item + for item in node.body + if (isinstance(item, ast.FunctionDef) and not item.name.startswith("_")) + or ( + isinstance(item, ast.AnnAssign) + and isinstance(item.target, ast.Name) + and not item.target.id.startswith("_") + ) + } + + +def stub_signature(node): + """Compare Python-callable names, kinds and defaults, ignoring annotations/self.""" + args = node.args + positional = [*args.posonlyargs, *args.args] + defaults = [inspect.Parameter.empty] * (len(positional) - len(args.defaults)) + defaults += [ast.literal_eval(value) for value in args.defaults] + parameters = [] + for index, (arg, default) in enumerate(zip(positional, defaults, strict=True)): + if arg.arg == "self": + continue + kind = ( + inspect.Parameter.POSITIONAL_ONLY + if index < len(args.posonlyargs) + else inspect.Parameter.POSITIONAL_OR_KEYWORD + ) + parameters.append(inspect.Parameter(arg.arg, kind, default=default)) + if args.vararg: + parameters.append( + inspect.Parameter(args.vararg.arg, inspect.Parameter.VAR_POSITIONAL) + ) + for arg, default in zip(args.kwonlyargs, args.kw_defaults, strict=True): + parameters.append( + inspect.Parameter( + arg.arg, + inspect.Parameter.KEYWORD_ONLY, + default=inspect.Parameter.empty + if default is None + else ast.literal_eval(default), + ) + ) + if args.kwarg: + parameters.append( + inspect.Parameter(args.kwarg.arg, inspect.Parameter.VAR_KEYWORD) + ) + return inspect.Signature(parameters) + + +def runtime_signature(value): + return inspect.Signature( + [ + parameter.replace(annotation=inspect.Parameter.empty) + for parameter in inspect.signature(value).parameters.values() + if parameter.name != "self" + ] + ) + + +def check_class(node, runtime): + members = public_members(node) + actual = {name for name in vars(runtime) if not name.startswith("_")} + assert actual == members.keys(), f"{node.name}: runtime/stub public members differ" + for name, member in members.items(): + value = getattr(runtime, name) + if isinstance(member, ast.AnnAssign): + assert isinstance(value, runtime), ( + f"{node.name}.{name}: expected enum member" + ) + elif any( + isinstance(d, ast.Name) and d.id == "property" + for d in member.decorator_list + ): + assert inspect.isgetsetdescriptor(value), ( + f"{node.name}.{name}: expected property" + ) + else: + assert runtime_signature(value) == stub_signature(member), ( + f"{node.name}.{name}" + ) + constructor = next( + ( + item + for item in node.body + if isinstance(item, ast.FunctionDef) and item.name == "__init__" + ), + None, + ) + if constructor is not None: + assert runtime_signature(runtime) == stub_signature(constructor), node.name + + +def test_public_stubs_match_compiled_extension(): + classes = stub_classes() + assert set(lance_csr.__all__) == classes.keys() + assert { + name for name in vars(_internal.graph) if not name.startswith("_") + } == classes.keys() + for name, node in classes.items(): + runtime = getattr(_internal.graph, name) + assert getattr(lance_csr, name) is runtime + check_class(node, runtime) + + +def test_stub_check_detects_missing_method(): + node = stub_classes()["CypherEngine"] + node.body = [ + item + for item in node.body + if not isinstance(item, ast.FunctionDef) or item.name != "refresh" + ] + with pytest.raises(AssertionError, match="public members differ"): + check_class(node, lance_csr.CypherEngine) + + +def test_stub_check_detects_changed_default(): + node = stub_classes()["CypherEngine"] + execute = next( + item + for item in node.body + if isinstance(item, ast.FunctionDef) and item.name == "execute" + ) + execute.args.defaults = [ast.Constant(False)] + with pytest.raises(AssertionError, match="CypherEngine.execute"): + check_class(node, lance_csr.CypherEngine) diff --git a/python/python/tests/typing/public_api.py b/python/python/tests/typing/public_api.py new file mode 100644 index 0000000..1edc328 --- /dev/null +++ b/python/python/tests/typing/public_api.py @@ -0,0 +1,95 @@ +"""Strict static consumer; checked by make lint, not executed by pytest.""" + +from datetime import date +from typing import assert_type + +from lance import LanceDataset +from lance_csr import ( + CypherEngine, + CypherQuery, + DirNamespace, + DistanceMetric, + ExecutionStrategy, + GraphConfig, + GraphConfigBuilder, + SqlDialect, + SqlEngine, + SqlQuery, + VectorSearch, + __version__, +) +from pyarrow import RecordBatch, RecordBatchReader, Table + + +def public_api( + table: Table, batch: RecordBatch, reader: RecordBatchReader, dataset: LanceDataset +) -> None: + assert_type(__version__, str) + builder = GraphConfig.builder().node_table("Person", "id") + assert_type(builder, GraphConfigBuilder) + builder = builder.rel_table("KNOWS", "Person", "Person", "src", "dst") + config = builder.build() + assert_type(config, GraphConfig) + assert_type(config.node_labels(), list[str]) + assert_type(config.relationship_types(), list[str]) + assert_type(GraphConfigBuilder().build(), GraphConfig) + assert_type(builder.with_node_label("City", "id"), GraphConfigBuilder) + assert_type(builder.with_relationship("R", "src", "dst"), GraphConfigBuilder) + + datasets = {"Person": table} + engine = CypherEngine(config, datasets, use_csr=True) + query = "MATCH (p:Person) WHERE p.id = $id RETURN p.id" + params = {"id": 1} + assert_type(engine.execute(query, params), Table) + assert_type(engine.execute_batches(query, params), RecordBatchReader) + assert_type(engine.explain(query, params), str) + assert_type(engine.refresh(), None) + assert_type(engine.config(), GraphConfig) + assert_type(engine.dataset_versions(), dict[str, int]) + assert_type(engine.node_index_stats()["person"]["rows"], int) + assert_type(engine.node_index_stats()["person"]["build_ms"], float) + assert_type(engine.csr_stats()["knows"]["edges"], int) + assert_type(engine.csr_stats()["knows"]["version"], int | None) + CypherEngine(config, {"Person": batch}) + CypherEngine(config, {"Person": reader}) + CypherEngine(config, {"Person": dataset}) + + one = CypherQuery(query).with_config(config).with_parameter("id", 1) + assert_type(one, CypherQuery) + one.with_parameter("day", date(2026, 1, 1)) + assert_type(one.query_text(), str) + assert_type(one.parameters(), dict[str, int | float | str | bool | None]) + assert_type(one.variables(), list[str]) + assert_type(one.node_labels(), list[str]) + assert_type(one.relationship_types(), list[str]) + assert_type(one.execute(datasets, ExecutionStrategy.DataFusion), Table) + assert_type(one.explain(datasets), str) + assert_type(one.to_sql(datasets, SqlDialect.PostgreSql), str) + namespace = DirNamespace("/tmp/graph") + assert_type(namespace.base_uri, str) + assert_type(one.execute_with_namespace(namespace), Table) + + vector = ( + VectorSearch("p.embedding") + .query_vector([1.0, 0.0]) + .metric(DistanceMetric.Cosine) + .top_k(5) + .include_distance(True) + .distance_column_name("distance") + .use_lance_index(False) + ) + assert_type(vector, VectorSearch) + assert_type(vector.search(table), Table) + assert_type(one.execute_with_vector_rerank(datasets, vector), Table) + assert_type(engine.execute_with_vector_rerank(query, vector), Table) + sql = SqlQuery("SELECT id FROM Person") + assert_type(sql.sql(), str) + assert_type(sql.execute(datasets), Table) + assert_type(sql.explain(datasets), str) + assert_type(SqlEngine(datasets).execute(sql.sql()), Table) + + # Rejected calls are checked too: removing a type restriction makes these unused. + engine.execute(query, {"id": object()}) # pyright: ignore[reportArgumentType] + vector.top_k("5") # pyright: ignore[reportArgumentType] + one.with_parameter("ids", [1, 2]) # pyright: ignore[reportArgumentType] + namespace.base_uri = "elsewhere" # pyright: ignore[reportAttributeAccessIssue] diff --git a/python/uv.lock b/python/uv.lock index 26fda10..5a277bf 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -51,6 +51,7 @@ dependencies = [ dev = [ { name = "maturin", extra = ["patchelf"] }, { name = "pandas" }, + { name = "pyarrow-stubs" }, { name = "pyright" }, { name = "pytest" }, { name = "ruff" }, @@ -66,6 +67,7 @@ requires-dist = [ dev = [ { name = "maturin", extras = ["patchelf"], specifier = ">=1.9,<2" }, { name = "pandas" }, + { name = "pyarrow-stubs", specifier = ">=20.0.0.20260819" }, { name = "pyright" }, { name = "pytest" }, { name = "ruff" }, @@ -352,6 +354,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/96/be/7b81a44d6a8e70581dcc1d6f01541f9000a973b1e5d75394aec91e7b179a/pyarrow-26.0.0-cp315-cp315t-win_amd64.whl", hash = "sha256:68cd662e9e2b00876a131950cf32336ace2d0865e1f9418763e3d3be8481dfa4", size = 29389953, upload-time = "2026-10-09T08:26:18.277Z" }, ] +[[package]] +name = "pyarrow-stubs" +version = "20.0.0.20260819" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyarrow" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/12/a7/8a2ca91ffe4c6576207f932de655d7d8a36485c520dccce70ce7d492b256/pyarrow_stubs-20.0.0.20260819.tar.gz", hash = "sha256:150710a72248bc834bf048d3092713f070904a4af76d40289c43afb3ee189823", size = 238222, upload-time = "2026-08-19T05:52:53.618Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/65/6c/eea1d03e475217aea95b1d52aee09c97575d05bbc592c39c085b71dab89f/pyarrow_stubs-20.0.0.20260819-py3-none-any.whl", hash = "sha256:297e60b6e5314739c082b4757d090d8be6047465510eb0684ca954ef7ea58be3", size = 235949, upload-time = "2026-08-19T05:52:54.711Z" }, +] + [[package]] name = "pydantic" version = "2.14.1"