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 b528333..8ffd132 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -19,6 +19,35 @@ code. The guides link here for exact signatures and link back for examples. This page covers the Python API. For Rust, see [Using from Rust](rust.md) and the [API documentation on docs.rs](https://docs.rs/lance-csr). +## 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 da2bc69..fa10669 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 6715e47..36acc3c 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 9d4d52d..2ccbec6 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 466cf6e..29a88c4 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"