Prebuilt rtree abi3 prototype

.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

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

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=[])