diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f73cf42..258955b 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,89 @@ 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 + musllinux x86_64 + - ubuntu-24.04-arm # manylinux + musllinux aarch64 + - windows-latest # win_amd64 + win32 + - 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 + + # Everything else builds its own wheel; this is the only job that exercises + # the artifacts we would actually publish, and the only one that checks the + # stable-ABI claim -- one cp311 wheel running on every supported CPython. + 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 the wheel directly: cibuildwheel emits a manylinux and a musllinux + # wheel, so a bare *.whl glob is ambiguous, while installing by package + # name would resolve `spatial-graph` from PyPI instead of this artifact. + - name: Install the built wheel with its real dependencies + run: | + uv venv + uv pip install wheelhouse/*manylinux*.whl pytest + - name: Test it + # test_prebuilt.py asserts the real property directly: constructing a + # prebuilt variant must not reach witty.compile_cython + run: uv run --no-sync pytest tests/test_prebuilt.py -v + deploy: name: Deploy - needs: test + needs: [test, build-wheels, build-sdist, test-abi3-wheel] if: success() && startsWith(github.ref, 'refs/tags/') && github.event_name != 'schedule' runs-on: ubuntu-latest @@ -85,16 +164,15 @@ jobs: 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/README.md b/README.md index 95174f2..3508bc4 100644 --- a/README.md +++ b/README.md @@ -26,7 +26,8 @@ where each node has an nD position (in time or space). * attribute access * minimal memory footprint * minimal dependencies - * `cython` / `witty` / `cheetah3` for runtime compilation + * `cython` / `witty` / `cheetah3`, used only when something has to be + compiled at runtime (see Cross-Platform Support) * numpy for array interfaces * PYX API for graph algorithms in C/C++ @@ -116,16 +117,34 @@ A `SpatialGraph` consists of three data structures: ## Cross-Platform Support -`spatial_graph` compiles C/C++ code at runtime, and as such needs access to a -compiler. If you already have one, great! You can use the PyPI package. +`spatial_graph` generates specialized C/C++ for the exact data types you ask +for. Where those types can be known in advance we compile them ahead of time +and ship them in the wheels; everything else is compiled on your machine the +first time it is used, which needs a C compiler. -If you (or your users) don't have a compiler installed, you either need to +**No compiler needed.** The PyPI wheels contain prebuilt `PointRTree` and +`LineRTree` variants for the common combinations: `float32`/`float64` +coordinates, 2 to 5 dimensions, and `int64`/`uint64` items -- as `int64[2]` / +`uint64[2]` for `LineRTree`, whose items are node pairs. If your R-tree matches +one of those -- as most do -- nothing is compiled, on any supported Python. -1. Install a compiler. This might be weird for non-technical users. -2. Install `spatial_graph` from `conda-forge`, where we include a compiler - (`clang`) in its dependencies. +**Compiler needed.** Two cases fall back to compiling at runtime: -### Why is this so complicated? +1. `Graph`, `DiGraph`, `SpatialGraph` and `SpatialDiGraph`. Their node and edge + attribute types are only known when you construct the graph, so they cannot + be enumerated ahead of time. +2. R-trees outside the prebuilt set above (an `int32` item type, say, or 6 + dimensions). + +If you or your users need those without a compiler, you can still install +`spatial_graph` from `conda-forge`, where we include a compiler (`clang`) in +its dependencies. + +The wheels are `abi3` (stable ABI) and require Python 3.11 or newer, so one +wheel per platform covers every supported CPython. Python 3.10 users should +pin to a release before this one. + +### Why can't everything be prebuilt? There is no cross-platform C/C++ compiler that we can install using `pip`. [`numba`](https://github.com/numba/numba) is maybe the closest to having solved diff --git a/pyproject.toml b/pyproject.toml index de66160..62d786b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,17 +1,34 @@ [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. `include-package-data` (on by default here) would pick them up from +# the VCS/sdist manifest anyway; listing them keeps that explicit and works +# even when no file finder is available. Patterns are package-relative, so the +# vendored C under `_rtree/src/` and `_graph/src/` needs the `src/` prefix. +"*" = ["py.typed", "*.pyx", "src/*.c", "src/*.h", "src/LICENSE*", "src/*.md"] [project] name = "spatial-graph" dynamic = ["version"] description = "A spatial graph datastructure for python." readme = "README.md" -requires-python = ">=3.10" -license = { text = "MIT" } +requires-python = ">=3.11" +license = "MIT" +license-files = ["LICENSE"] authors = [ { email = "funkej@janelia.hhmi.org", name = "Jan Funke" }, { email = "talley.lambert@gmail.com", name = "Talley Lambert" }, @@ -19,8 +36,6 @@ authors = [ 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 +48,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 +77,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..f816e3c --- /dev/null +++ b/setup.py @@ -0,0 +1,188 @@ +"""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 the two paths are +generated from one source. One wheel per platform then covers every supported +CPython, and those variants are never compiled on the user's machine -- so no +C compiler is invoked for them. + +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 _src_on_path() -> None: + """Make the package under `src/` importable, for the helpers shared with it. + + Callers import `spatial_graph` only once they know they need it: importing + it at module scope would make every command -- `sdist` and `egg_info` + included -- fail if anything the package imports is missing from the build + environment. + """ + if str(SRC) not in sys.path: + sys.path.insert(0, str(SRC)) + + +def prebuilt_extensions() -> list[Extension]: + """Render every prebuilt RTree variant and declare it as an extension.""" + from Cython.Build import cythonize + + _src_on_path() + 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) + + # the template `#include`s these at compile time, so they never appear in + # `sources`; without `depends` an edit to the C core would leave every + # prebuilt module stale (the JIT path declares the same set to witty) + rtree_dir = SRC / "spatial_graph" / "_rtree" + depends = [str(rtree_dir / "src" / f) for f in ("rtree.c", "rtree.h", "config.h")] + + 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)], + depends=depends, + include_dirs=[str(rtree_dir)], + 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. + + Includes `Python.h` so that a box with a C compiler but no development + headers -- the usual shape of this failure -- is caught here rather than + part-way through the real build. + """ + try: + from distutils.ccompiler import new_compiler + from distutils.sysconfig import customize_compiler, get_python_inc + + 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("#include \nint main(void) { return 0; }\n") + # mirror build_ext, which passes both include dirs when the + # platform-specific one differs (that is where pyconfig.h can live) + includes = [get_python_inc()] + if (plat_inc := get_python_inc(plat_specific=True)) not in includes: + includes.append(plat_inc) + compiler.compile([str(probe)], output_dir=tmp, include_dirs=includes) + except Exception: + return False + return True + + +# Rendering and cythonizing every variant is wasted work for commands that only +# want metadata -- without this, `build --sdist` pays for a full codegen pass. +METADATA_ONLY = {"egg_info", "dist_info", "sdist"} +NEEDS_EXTENSIONS = { + "bdist_egg", + "bdist_wheel", + "build", + "build_ext", + "build_py", + "develop", + "editable_wheel", + "install", +} + + +def metadata_only() -> bool: + """Whether this invocation asks for nothing that needs the extensions. + + Matches on recognized command names only, so option values (`--dist-dir + /tmp/x`) are ignored, and anything unrecognized falls through to building + -- skipping wrongly would silently yield a wheel with no prebuilt modules. + """ + seen = {arg for arg in sys.argv[1:] if arg in METADATA_ONLY | NEEDS_EXTENSIONS} + return bool(seen) and seen <= METADATA_ONLY + + +def should_prebuild() -> bool: + """Whether to compile prebuilt variants into this wheel.""" + if metadata_only(): + return False + + _src_on_path() + from spatial_graph._rtree._naming import env_enabled + + if env_enabled("SPATIAL_GRAPH_NO_PREBUILT"): + return False + if env_enabled("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 __name__ == "__main__": + 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/_graph/graph_base.py b/src/spatial_graph/_graph/graph_base.py index 52ac4cf..fc4ca90 100644 --- a/src/spatial_graph/_graph/graph_base.py +++ b/src/spatial_graph/_graph/graph_base.py @@ -4,9 +4,6 @@ from pathlib import Path from typing import TYPE_CHECKING, Any -import witty -from Cheetah.Template import Template - from spatial_graph._dtypes import DType from .views import EdgeAttrs, NodeAttrs @@ -75,6 +72,10 @@ def _build_wrapper( if not all(str.isidentifier(name) for name in edge_attr_dtypes): raise ValueError("Edge attribute names must be valid identifiers") + # imported here, not at module scope: it is only needed when something has + # to be compiled, and keeps `import spatial_graph` off Cheetah/Cython + from Cheetah.Template import Template + wrapper_template = Template( file=str(SRC_DIR / "wrapper_template.pyx"), compilerSettings={"directiveStartToken": "%"}, @@ -99,6 +100,11 @@ def _compile_graph( edge_attr_dtypes: Mapping[str, str] | None = None, directed: bool = False, ) -> type: + # graph attribute dtypes are only known at runtime, so this half can never + # be prebuilt; witty is imported here so that merely importing spatial_graph + # does not pull in the compilation toolchain + import witty + wrapper_template = _build_wrapper( node_dtype=node_dtype, node_attr_dtypes=node_attr_dtypes, diff --git a/src/spatial_graph/_rtree/_codegen.py b/src/spatial_graph/_rtree/_codegen.py new file mode 100644 index 0000000..5e09257 --- /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`: Cheetah and witty stay installed either way, but +a tree that resolves to a prebuilt module never invokes them, and so never +needs a C compiler. +""" + +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. Both tree classes are +# public API and usable on their own, so both are prebuilt: it is what lets an +# rtree-only user install without a C compiler. Trimming the matrix means +# dropping entries below; `SPATIAL_GRAPH_NO_PREBUILT=1` skips prebuilding +# entirely, for machines that cannot compile at build time. +ITEM_BASES = ("int64", "uint64") +COORD_DTYPES = ("float32", "float64") +DIMS = (2, 3, 4, 5) + + +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) + 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..3b06a99 --- /dev/null +++ b/src/spatial_graph/_rtree/_naming.py @@ -0,0 +1,63 @@ +"""Naming and switches 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`, and importing Cheetah or witty here would pull the compilation +toolchain into the common case, which is the whole thing prebuilding avoids. +""" + +from __future__ import annotations + +import hashlib +import os +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" + +_FALSEY = {"", "0", "false", "no", "off"} + + +def env_enabled(name: str) -> bool: + """Whether the named on/off environment variable is switched on. + + `SPATIAL_GRAPH_NO_PREBUILT=0` should mean "no, don't skip prebuilding"; + plain truthiness would read it as "yes", since any non-empty string is + true. Both the build and the runtime lookup go through here so they can + never read the same variable differently. + """ + return os.environ.get(name, "").strip().lower() not in _FALSEY + + +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. It deliberately does not cover the + template itself, the vendored C, or the compiler flags: keeping those in + step is the build's job (see `depends` in `setup.py`). + """ + 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..8c79055 100644 --- a/src/spatial_graph/_rtree/rtree.py +++ b/src/spatial_graph/_rtree/rtree.py @@ -1,15 +1,16 @@ from __future__ import annotations +import importlib 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, env_enabled, 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 +20,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 env_enabled("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 rather than at module scope so that trees served from a + prebuilt module never invoke them -- and so never need a C compiler. + """ + import witty + + from ._codegen import build_wrapper + + wrapper = build_wrapper(cls, item_dtype, coord_dtype, dims) module = witty.compile_cython( wrapper, depends_on=[ @@ -62,9 +64,20 @@ 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. + """A generic RTree implementation, specialized for the given types. + + Common type combinations ship precompiled in the wheels; anything else is + compiled on the fly during instantiation, which requires a C compiler. Args: diff --git a/tests/test_prebuilt.py b/tests/test_prebuilt.py new file mode 100644 index 0000000..db95b37 --- /dev/null +++ b/tests/test_prebuilt.py @@ -0,0 +1,97 @@ +"""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 numpy as np +import pytest +import witty + +from spatial_graph import PointRTree +from spatial_graph._rtree._codegen import iter_specs +from spatial_graph._rtree._naming import env_enabled, 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 env_enabled("SPATIAL_GRAPH_REQUIRE_PREBUILT"): + pytest.skip("SPATIAL_GRAPH_REQUIRE_PREBUILT not enabled") + assert has_prebuilt, "install shipped no prebuilt rtree modules" + + +@requires_prebuilt +@pytest.mark.parametrize("spec", list(iter_specs()), ids=str) +def test_declared_specs_are_shipped_and_never_compile(spec, monkeypatch): + """Every declared variant must construct without invoking the compiler. + + This is the property prebuilding exists for. Blocking `compile_cython` + asserts it directly, in the environment users actually get -- witty is an + unconditional dependency, so it is always installed; what must not happen + is that it gets *used*. + """ + + def no_compiling(*args, **kwargs): + raise AssertionError(f"{spec} triggered runtime compilation") + + monkeypatch.setattr(witty, "compile_cython", no_compiling) + tree = spec.cls(spec.item_dtype, spec.coord_dtype, spec.dims) + assert "_prebuilt" in type(tree._ctree).__module__ + + +@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