setup.py

"""
Setup script for sqlite-vec
Builds the C extension and bundles Python helper functions
"""

import os
import subprocess
from setuptools import setup, Extension
from setuptools.command.build_ext import build_ext
from setuptools.command.build_py import build_py

# Create sqlite_vec package directory with placeholder __init__.py
# This will be overwritten during build with the real content
os.makedirs("sqlite_vec", exist_ok=True)
if not os.path.exists("sqlite_vec/__init__.py"):
    with open("sqlite_vec/__init__.py", "w") as f:
        f.write("# Placeholder - will be generated during build\n")


class MakeBuild(build_ext):
    """Custom build that uses Make to compile the extension"""

    def run(self):
        import shutil

        # Run vendor script to get dependencies
        subprocess.check_call(["./scripts/vendor.sh"], shell=True)

        # Build loadable extension
        subprocess.check_call(["make", "loadable"])

        # Determine extension suffix
        ext_suffix = ".so"
        if os.name == "nt":
            ext_suffix = ".dll"
        elif os.uname().sysname == "Darwin":
            ext_suffix = ".dylib"

        # Read VERSION file
        version = open("VERSION").read().strip()

        # Read helper functions
        helper_code = open("bindings/python/extra_init.py").read()

        # Create a Python package that matches upstream API
        with open("sqlite_vec/__init__.py", "w") as f:
            f.write(f'''
from os import path
import sqlite3

__version__ = "{version}"
__version_info__ = tuple(__version__.split("."))

def loadable_path():
  """ Returns the full path to the sqlite-vec loadable SQLite extension bundled with this package """

  loadable_path = path.join(path.dirname(__file__), "vec0")
  return path.normpath(loadable_path)

def load(conn: sqlite3.Connection)  -> None:
  """ Load the sqlite-vec SQLite extension into the given database connection. """

  conn.load_extension(loadable_path())

''')
            f.write(helper_code)

        # Copy built extension to package directory
        src = os.path.join("dist", f"vec0{ext_suffix}")
        dst = os.path.join("sqlite_vec", f"vec0{ext_suffix}")
        shutil.copy2(src, dst)

    def get_outputs(self):
        # Return the path to the built extension
        ext_suffix = ".so"
        if os.name == "nt":
            ext_suffix = ".dll"
        elif os.uname().sysname == "Darwin":
            ext_suffix = ".dylib"
        return [os.path.join("sqlite_vec", f"vec0{ext_suffix}")]


class CustomBuildPy(build_py):
    """Ensure build_ext runs before build_py to generate files"""

    def run(self):
        self.run_command("build_ext")
        super().run()


setup(
    cmdclass={"build_py": CustomBuildPy, "build_ext": MakeBuild},
    ext_modules=[Extension("vec0", sources=[])],
)