feat(python): ship public API type stubs

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()

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

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

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"

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. "

python/python/lance_csr/_internal/__init__.pyi

@@ -0,0 +1,3 @@
+from . import graph as graph
+
+__version__: str

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: ...