diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f73cf42..40e40e4 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -29,17 +29,19 @@ jobs: matrix: # ubuntu: full python range x both resolutions os: [ubuntu-latest] - python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] + python-version: ["3.11", "3.12", "3.13", "3.14"] resolution: [lowest-direct, highest] # windows/macos: only the endpoints, highest resolution include: - - { os: windows-latest, python-version: "3.10", resolution: highest } + - { os: windows-latest, python-version: "3.11", resolution: highest } - { os: windows-latest, python-version: "3.14", resolution: highest } - - { os: macos-latest, python-version: "3.10", resolution: highest } + - { os: macos-latest, python-version: "3.11", resolution: highest } - { os: macos-latest, python-version: "3.14", resolution: highest } env: UV_RESOLUTION: ${{ matrix.resolution }} + # a build that silently falls back to pure-Python must fail, not go green + SPATIAL_GRAPH_REQUIRE_PREBUILT: "1" steps: - uses: actions/checkout@v4 @@ -48,8 +50,12 @@ jobs: python-version: ${{ matrix.python-version }} enable-cache: true cache-dependency-glob: "**/pyproject.toml" + # --no-editable so we test the built wheel, prebuilt rtree modules and all, + # rather than an editable install of src/ + - name: Install as a built wheel + run: uv sync --no-dev --group test --no-editable - name: Test with coverage - run: uv run --no-dev --group test pytest -v --cov=spatial_graph --cov-report=xml + run: uv run --no-sync pytest -v --cov=spatial_graph --cov-report=xml - uses: codecov/codecov-action@v5 with: token: ${{ secrets.CODECOV_TOKEN }} @@ -66,16 +72,104 @@ jobs: enable-cache: true - name: install - run: uv sync --no-dev --group test-codspeed + run: uv sync --no-dev --group test-codspeed --no-editable - name: Run benchmarks uses: CodSpeedHQ/action@v3 with: run: uv run pytest -W ignore --codspeed -v --color=yes + # One abi3 wheel per platform, covering every supported CPython. Also the only + # thing that produces PyPI-acceptable manylinux tags -- `uv build` alone emits + # `linux_x86_64`, which PyPI rejects. + build-wheels: + name: Wheels ${{ matrix.os }} + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: + - ubuntu-latest # manylinux x86_64 + - ubuntu-24.04-arm # manylinux aarch64 + - windows-latest # win_amd64 + - macos-15-intel # macOS x86_64 + - macos-latest # macOS arm64 + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 # setuptools-scm needs the tags + - uses: pypa/cibuildwheel@v4.1.1 + - uses: actions/upload-artifact@v4 + with: + name: wheels-${{ matrix.os }} + path: wheelhouse/*.whl + + build-sdist: + name: Sdist + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + - uses: astral-sh/setup-uv@v6 + - run: uv build --sdist + - uses: actions/upload-artifact@v4 + with: + name: sdist + path: dist/*.tar.gz + + # The claim this whole design rests on: one cp311-abi3 wheel runs on every + # supported CPython, with no compiler and no witty. + test-abi3-wheel: + name: abi3 wheel on py${{ matrix.python-version }} + needs: build-wheels + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.11", "3.12", "3.13", "3.14"] + env: + SPATIAL_GRAPH_REQUIRE_PREBUILT: "1" + steps: + - uses: actions/checkout@v4 + - uses: actions/download-artifact@v4 + with: + name: wheels-ubuntu-latest + path: wheelhouse + - uses: astral-sh/setup-uv@v6 + with: + python-version: ${{ matrix.python-version }} + - name: Install the wheel with numpy alone + run: | + uv venv + uv pip install numpy + uv pip install --no-deps wheelhouse/*.whl + - name: Prebuilt rtrees must work without witty, Cheetah or a compiler + run: | + uv run --no-sync python -c " + import sys, numpy as np + try: + import witty; sys.exit('witty present; test is not conclusive') + except ImportError: pass + from spatial_graph import PointRTree + t = PointRTree('int64', 'float32', 3) + t.insert_point_items(np.array([1, 2], dtype='int64'), + np.ascontiguousarray([[0,0,0],[9,9,9]], dtype='float32')) + mod = type(t._ctree).__module__ + assert '_prebuilt' in mod, mod + found = t.search(np.array([0,0,0],'float32'), np.array([1,1,1],'float32')) + assert found.ravel().tolist() == [1], found + print('ok:', mod)" + # the test module imports the codegen (and so Cheetah), so pull the real + # dependency set back in before running the suite + - name: Run the prebuilt test suite against the wheel + run: | + uv pip install wheelhouse/*.whl pytest + uv run --no-sync pytest tests/test_prebuilt.py -v + deploy: name: Deploy - needs: test + needs: [test, test-abi3-wheel, build-sdist] if: success() && startsWith(github.ref, 'refs/tags/') && github.event_name != 'schedule' runs-on: ubuntu-latest @@ -84,17 +178,15 @@ jobs: contents: write steps: - - uses: actions/checkout@v4 + - uses: actions/download-artifact@v4 with: - fetch-depth: 0 - - uses: astral-sh/setup-uv@v6 + pattern: wheels-* + path: dist + merge-multiple: true + - uses: actions/download-artifact@v4 with: - python-version: ${{ matrix.python-version }} - enable-cache: true - cache-dependency-glob: "**/pyproject.toml" - - - name: 👷 Build - run: uv build + name: sdist + path: dist - name: 🚢 Publish to PyPI uses: pypa/gh-action-pypi-publish@release/v1 diff --git a/pyproject.toml b/pyproject.toml index de66160..0dda8d2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,16 +1,28 @@ [build-system] -requires = ["hatchling", "hatch-vcs"] -build-backend = "hatchling.build" +requires = [ + "setuptools>=77", + "setuptools-scm>=8", + "Cython>=3.1", + "CT3>=3.3.3", + "numpy", # imported (not linked) while rendering the wrappers +] +build-backend = "setuptools.build_meta" + +[tool.setuptools_scm] -[tool.hatch.version] -source = "vcs" +[tool.setuptools.packages.find] +where = ["src"] + +[tool.setuptools.package-data] +# the JIT fallback compiles from these at runtime, so they must ship in the wheel +"*" = ["py.typed", "*.pyx", "*.c", "*.h", "LICENSE*", "*.md"] [project] name = "spatial-graph" dynamic = ["version"] description = "A spatial graph datastructure for python." readme = "README.md" -requires-python = ">=3.10" +requires-python = ">=3.11" license = { text = "MIT" } authors = [ { email = "funkej@janelia.hhmi.org", name = "Jan Funke" }, @@ -20,7 +32,6 @@ classifiers = [ "Development Status :: 3 - Alpha", "Intended Audience :: Developers", "License :: OSI Approved :: MIT License", - "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", @@ -33,8 +44,7 @@ dependencies = [ "numpy>=2.3.2; python_version >= '3.14'", "numpy>=2.1.0; python_version >= '3.13'", "numpy>=1.26.0; python_version >= '3.12'", - "numpy>=1.23.2; python_version >= '3.11'", - "numpy>=1.21.2", + "numpy>=1.23.2", "setuptools>=75.8.0", "typing_extensions>=4.5.0", # witty<=0.3.1 imports it without declaring it ] @@ -63,8 +73,18 @@ docs = [ homepage = "https://github.com/funkelab/spatial_graph" repository = "https://github.com/funkelab/spatial_graph" +[tool.cibuildwheel] +# a single abi3 build per platform covers every supported CPython +build = "cp311-*" +# never let a wheel silently degrade to pure Python +environment = { SPATIAL_GRAPH_REQUIRE_PREBUILT = "1" } +test-groups = ["test"] +# these exercise the prebuilt modules in the repaired wheel without needing a +# compiler; cross-version and numpy-only checks live in the CI workflow +test-command = "pytest {project}/tests/test_prebuilt.py -q" + [tool.ruff] -target-version = "py310" +target-version = "py311" line-length = 88 fix = true unsafe-fixes = true diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..67c98d7 --- /dev/null +++ b/setup.py @@ -0,0 +1,125 @@ +"""Compile RTree variants ahead of time into a stable-ABI (abi3) wheel. + +Renders the same pyx wrappers the runtime would JIT-compile (via +`_rtree._codegen`) for every variant in `iter_specs()`, so prebuilt and +JIT-compiled modules can never disagree. One wheel per platform then covers +every supported CPython, and users never need a C compiler for those variants. + +Set `SPATIAL_GRAPH_NO_PREBUILT=1` to build a pure-Python wheel instead, or +`SPATIAL_GRAPH_REQUIRE_PREBUILT=1` (as CI does) to turn a failure to compile +into a hard error rather than a silent fall back to JIT. +""" + +from __future__ import annotations + +import os +import sys +import tempfile +import warnings +from pathlib import Path + +from setuptools import Extension, setup + +ROOT = Path(__file__).parent +SRC = ROOT / "src" +PREBUILT_PKG = "spatial_graph._rtree._prebuilt" + +# The wrappers pass numpy arrays as typed memoryviews, which compile to +# PyObject_GetBuffer/PyBuffer_Release. Those entered the limited API in 3.11 +# (moved from cpython/object.h, excluded under Py_LIMITED_API, to pybuffer.h), +# so 3.11 is the floor for a stable-ABI build -- and matches requires-python. +ABI3_MIN = (3, 11) +ABI3_TAG = f"cp{ABI3_MIN[0]}{ABI3_MIN[1]}" +ABI3_HEX = f"0x{ABI3_MIN[0]:02x}{ABI3_MIN[1]:02x}0000" + +WIN = sys.platform == "win32" + + +def prebuilt_extensions() -> list[Extension]: + """Render every prebuilt RTree variant and declare it as an extension.""" + from Cython.Build import cythonize + + sys.path.insert(0, str(SRC)) + from spatial_graph._rtree._codegen import build_wrapper, iter_specs + from spatial_graph._rtree._naming import module_name + + pyx_dir = ROOT / "build" / "prebuilt-pyx" + pyx_dir.mkdir(parents=True, exist_ok=True) + + extensions = [] + for spec in iter_specs(): + name = module_name(spec.cls, spec.item_dtype, spec.coord_dtype, spec.dims) + source = build_wrapper(spec.cls, spec.item_dtype, spec.coord_dtype, spec.dims) + path = pyx_dir / f"{name}.pyx" + # only rewrite when changed, so cythonize can skip unchanged variants + if not path.is_file() or path.read_text() != source: + path.write_text(source) + extensions.append( + Extension( + f"{PREBUILT_PKG}.{name}", + sources=[str(path)], + include_dirs=[str(SRC / "spatial_graph" / "_rtree")], + extra_compile_args=["/O2"] if WIN else ["-O3", "-Wno-unreachable-code"], + define_macros=[ + ("Py_LIMITED_API", ABI3_HEX), + *([("RTREE_NOATOMICS", "1")] if WIN else []), + ], + py_limited_api=True, + ) + ) + + return cythonize( + extensions, + language_level=3, + quiet=True, + nthreads=0 if WIN else os.cpu_count(), + ) + + +def can_compile() -> bool: + """Whether this machine can build a C extension at all.""" + from distutils.ccompiler import new_compiler + from distutils.sysconfig import customize_compiler + + compiler = new_compiler() + customize_compiler(compiler) # picks up CC/CFLAGS, as build_ext does + with tempfile.TemporaryDirectory() as tmp: + probe = Path(tmp, "probe.c") + probe.write_text("int main(void) { return 0; }\n") + try: + compiler.compile([str(probe)], output_dir=tmp) + except Exception: + return False + return True + + +def should_prebuild() -> bool: + """Whether to compile prebuilt variants into this wheel.""" + if os.getenv("SPATIAL_GRAPH_NO_PREBUILT"): + return False + if os.getenv("SPATIAL_GRAPH_REQUIRE_PREBUILT"): + return True # CI: never let a build silently degrade + if can_compile(): + return True + # Installing from an sdist without a compiler must keep working: fall back to + # a pure-Python wheel that JIT-compiles on first use, as it did before + # prebuilding existed. + warnings.warn( + "No usable C compiler found; building spatial-graph without prebuilt " + "rtree modules. A C compiler will be needed the first time an RTree is " + "used.", + stacklevel=1, + ) + return False + + +if should_prebuild(): + setup( + ext_modules=prebuilt_extensions(), + options={ + "bdist_wheel": {"py_limited_api": ABI3_TAG}, + "build_ext": {"parallel": os.cpu_count()}, + }, + ) +else: + setup(ext_modules=[]) diff --git a/src/spatial_graph/__init__.py b/src/spatial_graph/__init__.py index 4ec1dee..c1ab15a 100644 --- a/src/spatial_graph/__init__.py +++ b/src/spatial_graph/__init__.py @@ -1,4 +1,5 @@ from importlib.metadata import PackageNotFoundError, version +from typing import TYPE_CHECKING, Any try: __version__ = version("spatial_graph") @@ -6,10 +7,25 @@ __version__ = "unknown" -from ._graph import DiGraph, Graph, GraphBase from ._rtree import LineRTree, PointRTree -from ._spatial_graph import SpatialDiGraph, SpatialGraph, SpatialGraphBase -from ._util import create_graph + +if TYPE_CHECKING: + from ._graph import DiGraph, Graph, GraphBase + from ._spatial_graph import SpatialDiGraph, SpatialGraph, SpatialGraphBase + from ._util import create_graph + +# the graph half is always JIT-compiled, and importing it pulls in witty and +# Cheetah. Deferring it keeps `PointRTree`/`LineRTree` -- which ship prebuilt -- +# usable with numpy alone. +_LAZY = { + "DiGraph": "._graph", + "Graph": "._graph", + "GraphBase": "._graph", + "SpatialDiGraph": "._spatial_graph", + "SpatialGraph": "._spatial_graph", + "SpatialGraphBase": "._spatial_graph", + "create_graph": "._util", +} __all__ = [ "DiGraph", @@ -22,3 +38,15 @@ "SpatialGraphBase", "create_graph", ] + + +def __getattr__(name: str) -> Any: + if module := _LAZY.get(name): + import importlib + + return getattr(importlib.import_module(module, __name__), name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def __dir__() -> list[str]: + return __all__ diff --git a/src/spatial_graph/_rtree/_codegen.py b/src/spatial_graph/_rtree/_codegen.py new file mode 100644 index 0000000..3b30f35 --- /dev/null +++ b/src/spatial_graph/_rtree/_codegen.py @@ -0,0 +1,72 @@ +"""What RTree variants get prebuilt, and how their pyx wrappers are rendered. + +Used on the JIT path and by `setup.py`, so prebuilt and JIT-compiled modules are +always generated from the same source. Requires Cheetah, and is therefore +imported lazily by `rtree.py` -- installs that stay on the prebuilt path need +neither Cheetah nor witty. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING, NamedTuple + +from Cheetah.Template import Template + +from spatial_graph._dtypes import DType + +from .line_rtree import LineRTree +from .point_rtree import PointRTree + +if TYPE_CHECKING: + from collections.abc import Iterator + + from .rtree import RTree + +TEMPLATE = Path(__file__).parent / "wrapper_template.pyx" + +# Variants compiled ahead of time into binary wheels. `PointRTree` is what makes +# a compiler unnecessary for rtree-only users; `LineRTree` is only reached via +# `SpatialGraph`, whose graph half is JIT-compiled regardless, so prebuilding it +# saves first-use compile time rather than removing a requirement. +ITEM_BASES = ("int64", "uint64") +COORD_DTYPES = ("float32", "float64") +DIMS = (2, 3, 4, 5) +PREBUILT_LINE_TREES = True + + +class Spec(NamedTuple): + cls: type[RTree] + item_dtype: str + coord_dtype: str + dims: int + + +def iter_specs() -> Iterator[Spec]: + """Yield every RTree variant that should be compiled into a wheel.""" + for base in ITEM_BASES: + for coord in COORD_DTYPES: + for dims in DIMS: + yield Spec(PointRTree, base, coord, dims) + if PREBUILT_LINE_TREES: + yield Spec(LineRTree, f"{base}[2]", coord, dims) + + +def build_wrapper( + cls: type[RTree], item_dtype: str, coord_dtype: str, dims: int +) -> str: + """Render the pyx wrapper for the given tree parameters.""" + wrapper_template = Template( + file=str(TEMPLATE), + compilerSettings={"directiveStartToken": "%"}, + ) + wrapper_template.item_dtype = DType(item_dtype) + wrapper_template.coord_dtype = DType(coord_dtype) + wrapper_template.dims = dims + wrapper_template.c_distance_function = cls.c_distance_function + wrapper_template.pyx_item_t_declaration = cls.pyx_item_t_declaration + wrapper_template.c_item_t_declaration = cls.c_item_t_declaration + wrapper_template.c_converter_functions = cls.c_converter_functions + wrapper_template.c_equal_function = cls.c_equal_function + + return str(wrapper_template) diff --git a/src/spatial_graph/_rtree/_naming.py b/src/spatial_graph/_rtree/_naming.py new file mode 100644 index 0000000..dcf34eb --- /dev/null +++ b/src/spatial_graph/_rtree/_naming.py @@ -0,0 +1,46 @@ +"""Deterministic naming for prebuilt RTree extension modules. + +Shared by the runtime lookup and `setup.py`, so the two can never disagree. +Deliberately depends only on `_dtypes` -- it sits on the import path of every +`PointRTree`, including installs with neither Cheetah nor witty available. +""" + +from __future__ import annotations + +import hashlib +from typing import TYPE_CHECKING + +from spatial_graph._dtypes import DType + +if TYPE_CHECKING: + from .rtree import RTree + +# subpackage holding ahead-of-time compiled modules; empty in a source checkout +PREBUILT_PACKAGE = f"{__package__}._prebuilt" + + +def _c_name(dtype: DType) -> str: + """Canonical, identifier-safe name for a dtype ("int64", "float", "int64x2").""" + base = dtype.base_c_type.removesuffix("_t") + return f"{base}x{dtype.size}" if dtype.is_array else base + + +def module_name(cls: type[RTree], item_dtype: str, coord_dtype: str, dims: int) -> str: + """Deterministic module name for the given tree parameters. + + Dtypes are canonicalized (so `int` and `int64` agree) and spelled out for + readability. The trailing digest covers the C/pyx code `cls` injects into the + template, so a subclass with custom code can never be served a prebuilt + module compiled from different code. + """ + parts = ( + cls.pyx_item_t_declaration, + cls.c_item_t_declaration, + cls.c_converter_functions, + cls.c_equal_function, + cls.c_distance_function, + ) + digest = hashlib.sha256("\0".join(parts).encode()).hexdigest()[:8] + item = _c_name(DType(item_dtype)) + coord = _c_name(DType(coord_dtype)) + return f"rtree_{item}_{coord}_d{dims}_{digest}" diff --git a/src/spatial_graph/_rtree/_prebuilt/__init__.py b/src/spatial_graph/_rtree/_prebuilt/__init__.py new file mode 100644 index 0000000..28a9aa6 --- /dev/null +++ b/src/spatial_graph/_rtree/_prebuilt/__init__.py @@ -0,0 +1,5 @@ +"""Ahead-of-time compiled RTree modules, populated at build time by `setup.py`. + +Empty in a plain source checkout: `_load_prebuilt` then finds nothing and every +tree is JIT-compiled, exactly as before prebuilding existed. +""" diff --git a/src/spatial_graph/_rtree/rtree.py b/src/spatial_graph/_rtree/rtree.py index 7bf280b..8afb50e 100644 --- a/src/spatial_graph/_rtree/rtree.py +++ b/src/spatial_graph/_rtree/rtree.py @@ -1,15 +1,17 @@ from __future__ import annotations +import importlib +import os import sys from pathlib import Path from typing import ClassVar import numpy as np -import witty -from Cheetah.Template import Template from spatial_graph._dtypes import DType +from ._naming import PREBUILT_PACKAGE, module_name + DEFINE_MACROS = [("RTREE_NOATOMICS", "1")] if sys.platform == "win32" else [] if sys.platform == "win32": # pragma: no cover EXTRA_COMPILE_ARGS = ["/O2"] @@ -19,33 +21,34 @@ SRC_DIR = Path(__file__).parent -def _build_wrapper( +def _load_prebuilt( cls: type[RTree], item_dtype: str, coord_dtype: str, dims: int -) -> str: - ############################################ - # create wrapper from template and compile # - ############################################ - - wrapper_template = Template( - file=str(SRC_DIR / "wrapper_template.pyx"), - compilerSettings={"directiveStartToken": "%"}, - ) - wrapper_template.item_dtype = DType(item_dtype) - wrapper_template.coord_dtype = DType(coord_dtype) - wrapper_template.dims = dims - wrapper_template.c_distance_function = cls.c_distance_function - wrapper_template.pyx_item_t_declaration = cls.pyx_item_t_declaration - wrapper_template.c_item_t_declaration = cls.c_item_t_declaration - wrapper_template.c_converter_functions = cls.c_converter_functions - wrapper_template.c_equal_function = cls.c_equal_function - - return str(wrapper_template) +) -> type | None: + """Return the ahead-of-time compiled tree class, or None if not shipped.""" + if os.getenv("SPATIAL_GRAPH_NO_PREBUILT"): + return None + name = module_name(cls, item_dtype, coord_dtype, dims) + try: + module = importlib.import_module(f"{PREBUILT_PACKAGE}.{name}") + except ImportError: + return None + return module.RTree -def _compile_tree( +def _jit_compile_tree( cls: type[RTree], item_dtype: str, coord_dtype: str, dims: int ) -> type: - wrapper = _build_wrapper(cls, item_dtype, coord_dtype, dims) + """Compile a tree with the system C compiler. + + Only reached for dtype combinations not shipped prebuilt; Cheetah and witty + are imported here so neither is needed by installs that stay on the + prebuilt path. + """ + import witty + + from ._codegen import build_wrapper + + wrapper = build_wrapper(cls, item_dtype, coord_dtype, dims) module = witty.compile_cython( wrapper, depends_on=[ @@ -62,6 +65,15 @@ def _compile_tree( return module.RTree +def _compile_tree( + cls: type[RTree], item_dtype: str, coord_dtype: str, dims: int +) -> type: + tree_cls = _load_prebuilt(cls, item_dtype, coord_dtype, dims) + if tree_cls is None: + tree_cls = _jit_compile_tree(cls, item_dtype, coord_dtype, dims) + return tree_cls + + class RTree: """A generic RTree implementation, compiled on-the-fly during instantiation. diff --git a/tests/test_prebuilt.py b/tests/test_prebuilt.py new file mode 100644 index 0000000..5ffa2b3 --- /dev/null +++ b/tests/test_prebuilt.py @@ -0,0 +1,86 @@ +"""Tests for ahead-of-time compiled rtree modules. + +The `requires_prebuilt` tests only mean something against an install that +actually shipped them, and are skipped otherwise -- except when +`SPATIAL_GRAPH_REQUIRE_PREBUILT` is set (as CI does), where their absence is +the very regression we want to catch. +""" + +from __future__ import annotations + +import os + +import numpy as np +import pytest + +from spatial_graph import PointRTree +from spatial_graph._rtree._codegen import iter_specs +from spatial_graph._rtree._naming import module_name +from spatial_graph._rtree.rtree import _load_prebuilt + +# the `_prebuilt` package always exists but is empty in a source checkout, so +# probe for a real module rather than for the package +has_prebuilt = _load_prebuilt(PointRTree, "int64", "float32", 2) is not None +requires_prebuilt = pytest.mark.skipif( + not has_prebuilt, reason="no prebuilt modules in this install" +) + + +def test_prebuilt_modules_were_shipped(): + """Guard against a wheel that silently degraded to pure Python.""" + if not os.getenv("SPATIAL_GRAPH_REQUIRE_PREBUILT"): + pytest.skip("SPATIAL_GRAPH_REQUIRE_PREBUILT not set") + assert has_prebuilt, "install shipped no prebuilt rtree modules" + + +@requires_prebuilt +@pytest.mark.parametrize("spec", list(iter_specs()), ids=str) +def test_every_declared_spec_is_shipped(spec): + """Every variant in `iter_specs` must actually resolve to a prebuilt module.""" + assert _load_prebuilt(spec.cls, spec.item_dtype, spec.coord_dtype, spec.dims) + + +@requires_prebuilt +def test_prebuilt_is_used_and_correct(): + tree = PointRTree("int64", "float32", 3) + assert "_prebuilt" in type(tree._ctree).__module__ + + items = np.array([10, 20, 30], dtype="int64") + points = np.ascontiguousarray([[0, 0, 0], [1, 1, 1], [9, 9, 9]], dtype="float32") + tree.insert_point_items(items, points) + + lo, hi = np.array([0, 0, 0], "float32"), np.array([2, 2, 2], "float32") + assert sorted(tree.search(lo, hi).ravel().tolist()) == [10, 20] + assert tree.nearest(np.array([8.9, 8.9, 8.9], "float32"), 1).ravel()[0] == 30 + + +def test_dtype_aliases_share_a_module(): + """`int`/`int64` and `float32`/`float` must not compile separate modules.""" + assert module_name(PointRTree, "int", "float32", 3) == module_name( + PointRTree, "int64", "float", 3 + ) + + +@pytest.mark.parametrize( + ("item_dtype", "coord_dtype", "dims"), + [("int32", "float32", 3), ("int64", "float32", 99)], +) +def test_unlisted_combination_falls_back_to_jit(item_dtype, coord_dtype, dims): + assert _load_prebuilt(PointRTree, item_dtype, coord_dtype, dims) is None + + +def test_subclass_with_custom_code_is_not_served_a_prebuilt_module(): + class CustomEquality(PointRTree): + c_equal_function = """ +inline bool equal(const item_t a, const item_t b) { return a == b; } +""" + + assert module_name(CustomEquality, "int64", "float32", 3) != module_name( + PointRTree, "int64", "float32", 3 + ) + assert _load_prebuilt(CustomEquality, "int64", "float32", 3) is None + + +def test_no_prebuilt_env_var_forces_jit(monkeypatch): + monkeypatch.setenv("SPATIAL_GRAPH_NO_PREBUILT", "1") + assert _load_prebuilt(PointRTree, "int64", "float32", 3) is None