setup.py

import os
import sys
import re
import subprocess
import sysconfig
import zipfile
from setuptools import setup, Extension
from setuptools.command.build_ext import build_ext
import setuptools
from distutils import log

log.set_verbosity(log.DEBUG)

# chdb wheel size optimization: use deflate level 9 instead of default level 6.
# This shaves ~0.3 MB off the wheel for free (no functional change).
_orig_zipfile_init = zipfile.ZipFile.__init__
def _chdb_zipfile_init(self, *args, **kwargs):
    if kwargs.get("compression", zipfile.ZIP_STORED) == zipfile.ZIP_DEFLATED and "compresslevel" not in kwargs:
        kwargs["compresslevel"] = 9
    _orig_zipfile_init(self, *args, **kwargs)
zipfile.ZipFile.__init__ = _chdb_zipfile_init


def get_python_ext_suffix():
    # Use Limited API suffix for cross-version compatibility
    return ".abi3.so"


# get the path of the current file
script_dir = os.path.dirname(os.path.abspath(__file__))
libdir = os.path.join(script_dir, "chdb")


def get_latest_git_tag(minor_ver_auto=False):
    try:
        # get latest tag commit
        completed_process = subprocess.run(
            ["git", "rev-list", "--tags", "--max-count=1"],
            capture_output=True,
            text=True,
        )
        if completed_process.returncode != 0:
            print(completed_process.stdout)
            print(completed_process.stderr)
            # get git version
            raise RuntimeError("Failed to get git latest tag commit ")
        output = completed_process.stdout.strip()
        # get latest tag name by commit
        completed_process = subprocess.run(
            ["git", "describe", "--tags", f"{output}"], capture_output=True, text=True
        )
        if completed_process.returncode != 0:
            print(completed_process.stdout)
            print(completed_process.stderr)
            # get git version
            raise RuntimeError("Failed to get git tag")
        output = completed_process.stdout.strip()
        # strip the v from the tag
        output = output[1:]
        parts = output.split(".")
        if len(parts) == 3:
            if minor_ver_auto:
                completed_process = subprocess.run(
                    ["git", "rev-list", "--count", f"v{output}..HEAD"],
                    capture_output=True,
                    text=True,
                )
                if completed_process.returncode != 0:
                    print(completed_process.stdout)
                    print(completed_process.stderr)
                    raise RuntimeError("Failed to get git rev-list")
                n = completed_process.stdout.strip()
                parts[2] = int(parts[2]) + int(n)
            return f"{parts[0]}.{parts[1]}.{parts[2]}"
    except Exception as e:
        print("Failed to get git tag. Error: ")
        print(e)
        raise


# When CHDB_LITE=1 in environment, rewrite pyproject.toml so the wheel is
# published as chdb-core-lite with its own PyPI-facing metadata (name,
# description, keywords). The long-description README is shared with chdb-core.
# Idempotent and safe to call multiple times.
def maybe_rewrite_pyproject_name_for_lite():
    if os.environ.get("CHDB_LITE", "0") != "1":
        return
    pyproject_file = os.path.join(script_dir, "pyproject.toml")
    with open(pyproject_file, "r") as f:
        content = f.read()
    new_content = re.sub(
        r'^name\s*=\s*"chdb-core"\s*$',
        'name = "chdb-core-lite"',
        content,
        flags=re.MULTILINE,
    )
    new_content = re.sub(
        r'^description\s*=\s*"[^"]*"\s*$',
        'description = "Lightweight build of chdb-core (in-process OLAP SQL engine powered by ClickHouse), trimmed for size-sensitive deployments such as serverless, containers and edge"',
        new_content,
        flags=re.MULTILINE,
    )
    new_content = re.sub(
        r'^keywords\s*=\s*\[[^\]]*\]\s*$',
        'keywords = ["chdb", "chdb-core", "chdb-core-lite", "lite", "clickhouse", "olap", "analytics", "database", "sql", "serverless", "edge"]',
        new_content,
        flags=re.MULTILINE,
    )
    if new_content != content:
        with open(pyproject_file, "w") as f:
            f.write(new_content)
        print("CHDB_LITE=1: rewrote pyproject.toml for chdb-core-lite (name, description, keywords)")


# Update version in pyproject.toml
def update_pyproject_version(version):
    pyproject_file = os.path.join(script_dir, "pyproject.toml")
    with open(pyproject_file, "r") as f:
        content = f.read()

    # Use regex to replace the version
    updated_content = re.sub(
        r'version\s*=\s*"[^"]*"', f'version = "{version}"', content
    )

    with open(pyproject_file, "w") as f:
        f.write(updated_content)


# As of Python 3.6, CCompiler has a `has_flag` method.
# cf http://bugs.python.org/issue26689
def has_flag(compiler, flagname):
    """Return a boolean indicating whether a flag name is supported on
    the specified compiler.
    """
    import tempfile

    with tempfile.NamedTemporaryFile("w", suffix=".cpp") as f:
        f.write("int main (int argc, char **argv) { return 0; }")
        try:
            compiler.compile([f.name], extra_postargs=[flagname])
        except setuptools.distutils.errors.CompileError:
            return False
    return True


def cpp_flag(compiler):
    """Return the -std=c++[11/2a] compiler flag.
    The c++2a is prefered over c++11 (when it is available).
    """
    if has_flag(compiler, "-std=c++2a"):
        return "-std=c++2a"
    elif has_flag(compiler, "-std=c++17"):
        return "-std=c++17"
    elif has_flag(compiler, "-std=c++14"):
        return "-std=c++14"
    elif has_flag(compiler, "-std=c++11"):
        return "-std=c++11"
    else:
        raise RuntimeError(
            "Unsupported compiler -- at least C++11 support " "is needed!"
        )


class BuildExt(build_ext):
    """A custom build extension for adding compiler-specific options."""

    def build_extensions(self):
        # add the _chdb.cpython-37m-darwin.so or _chdb.cpython-39-x86_64-linux.so to the chdb package
        self.distribution.package_data["chdb"] = [
            "chdb/_chdb" + get_python_ext_suffix()
        ]
        # super().build_extensions()


# this will be executed by python setup.py bdist_wheel
if __name__ == "__main__":
    try:
        # If CHDB_LITE=1, rewrite pyproject.toml name to chdb-core-lite
        # (must happen before setup() reads project metadata)
        maybe_rewrite_pyproject_name_for_lite()

        # get python extension file name
        chdb_so = libdir + "/_chdb" + get_python_ext_suffix()
        ext_modules = [
            Extension(
                "_chdb",
                sources=["programs/local/LocalChdb.cpp"],
                libraries=[],
                library_dirs=[libdir],
                extra_objects=[chdb_so],
                define_macros=[("Py_LIMITED_API", "0x03090000")],
                py_limited_api=True,
            ),
        ]
        # fix the version in chdb/__init__.py
        versionStr = get_latest_git_tag()
        # Call the function to update pyproject.toml
        # update_pyproject_version(versionStr)

        # scan the chdb directory and add all the .py and dynamic library files to the package
        pkg_files = []
        for root, dirs, files in os.walk(libdir):
            if "/build" in root or root.endswith("/build"):
                continue

            for file in files:
                if file.endswith(".py") or file.endswith(".pyi") or file == "py.typed":
                    pkg_files.append(os.path.join(root, file))
                # Include pybind11 nonlimitedapi libraries for all Python versions
                elif file.startswith("libpybind11nonlimitedapi_chdb_") and (
                    file.endswith(".dylib") or file.endswith(".so")
                ):
                    pkg_files.append(os.path.join(root, file))
                # Include pybind11 stub library
                elif file.startswith("libpybind11nonlimitedapi_stubs") and (
                    file.endswith(".dylib") or file.endswith(".so")
                ):
                    pkg_files.append(os.path.join(root, file))

        pkg_files.append(chdb_so)

        setup(
            packages=["chdb"],
            version=versionStr,
            include_package_data=False,
            package_data={"chdb": pkg_files},
            ext_modules=ext_modules,
            python_requires=">=3.9",
            install_requires=[
                "pyarrow>=13.0.0",
                "pandas>=2.1.0",
            ],
            cmdclass={"build_ext": BuildExt},
            test_suite="tests",
            zip_safe=False,
            options={
                "bdist_wheel": {
                    "py_limited_api": "cp39",
                }
            },
        )
    except Exception as e:
        print("Build from setup.py failed. Error: ")
        print(e)
        raise