setup.py

# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import glob
import os
import pathlib
import re
import shutil
import subprocess
import sys

import setuptools
import setuptools.command.build_ext


CC_FILES = glob.glob("src/**/*.cc", recursive=True)
H_FILES = glob.glob("src/**/*.h", recursive=True) + glob.glob("src/**/*.inl", recursive=True)
CMAKE_FILES = ['CMakeLists.txt', *glob.glob("file_lists/*", recursive=True)]
RELEVANT_SOURCE_FILES = sorted(CMAKE_FILES + CC_FILES + H_FILES)


class CMakeExtension(setuptools.Extension):
    def __init__(self, name, sourcedir, sources):
        setuptools.Extension.__init__(self, name, sources=sources)
        self.sourcedir = os.path.abspath(sourcedir)


class CMakeBuild(setuptools.command.build_ext.build_ext):
    def build_extension(self, ext):
        import ninja

        build_temp = os.path.join(self.build_temp, ext.name)
        os.makedirs(build_temp, exist_ok=True)
        build_env = {
            **os.environ,
            'CMAKE_LIBRARY_OUTPUT_DIRECTORY': '',
            'CMAKE_FORCE_PYBIND_CHROMOBIUS': '1',
        }

        osx_cmake_flags = []
        if sys.platform.startswith("darwin"):
            # Cross-compile support for macOS - respect ARCHFLAGS if set
            archs = re.findall(r"-arch (\S+)", os.environ.get("ARCHFLAGS", ""))
            if archs:
                osx_cmake_flags = ["-DCMAKE_OSX_ARCHITECTURES={}".format(";".join(archs))]
            else:
                import platform
                arch = platform.machine()
                if arch:
                    osx_cmake_flags = [f"-DCMAKE_OSX_ARCHITECTURES={arch}"]

        subprocess.check_call([
            "cmake",
            ext.sourcedir,
            f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={pathlib.Path(self.get_ext_fullpath(ext.name)).parent.absolute()}",
            f"-DPYTHON_EXECUTABLE={sys.executable}",
            f"-DCHROMOBIUS_VERSION_INFO={__version__}",
            *osx_cmake_flags,
            *[
                env_arg_item
                for env_arg_item in os.environ.get("CMAKE_ARGS", "").split(" ")
                if env_arg_item
            ],
            "-GNinja",
            f"-DCMAKE_MAKE_PROGRAM:FILEPATH={os.path.join(ninja.BIN_DIR, 'ninja')}",
        ], cwd=build_temp, env=build_env)

        subprocess.check_call([
            "cmake",
            "--build",
            ".",
            "--target",
            "chromobius_pybind",
        ], cwd=build_temp, env=build_env)


__version__ = '1.2.dev0'

with open("README.md", encoding="utf-8") as f:
    long_description = f.read()

# HACK: Workaround difficulties collecting data files for the package by just making a directory.
package_data_dir = pathlib.Path(__file__).parent / 'package_data'
package_data_dir.mkdir(exist_ok=True)
doc_chromobius = pathlib.Path(__file__).parent / 'doc' / 'chromobius.pyi'
if doc_chromobius.exists():
    shutil.copyfile(
        pathlib.Path(__file__).parent / 'doc' / 'chromobius.pyi',
        package_data_dir / 'chromobius.pyi',
    )

setuptools.setup(
    name="chromobius",
    version=__version__,
    author="Craig Gidney",
    url="https://github.com/quantumlib/chromobius",
    description="A fast implementation of the Möbius color code decoder.",
    long_description=long_description,
    long_description_content_type='text/markdown',
    maintainer="Google Quantum AI",
    maintainer_email="quantum-oss-maintainers@google.com",
    license="Apache-2.0",
    ext_modules=[CMakeExtension("chromobius", sourcedir=".", sources=RELEVANT_SOURCE_FILES)],
    cmdclass={"build_ext": CMakeBuild},
    python_requires=">=3.10",
    setup_requires=['ninja', 'pybind11~=2.11.1', 'cmake>=3.13'],
    install_requires=['numpy', 'stim'],

    # Needed on Windows to avoid the default `build` colliding with Bazel's `BUILD`.
    # Also, the replacement name is short to avoid blowing the 256 character path limit on windows.
    options={'build': {'build_base': 'b'}},

    # Add files in package_data_dir to the wheel.
    # I don't know why it has to be so esoteric, but I tried for hours.
    # This is the best I could come up with.
    packages=['chromobius'],
    package_dir={'chromobius': package_data_dir.name},
    package_data={'chromobius': [str(e) for e in package_data_dir.iterdir()]},
    include_package_data=True,
    classifiers=[
        "Operating System :: MacOS :: MacOS X",
        "Operating System :: Microsoft :: Windows",
        "Operating System :: POSIX :: Linux",
        "Programming Language :: Python :: 3",
        "Programming Language :: Python :: 3.10",
        "Programming Language :: Python :: 3.11",
        "Programming Language :: Python :: 3.12",
        "Programming Language :: Python :: 3.13",
        "Programming Language :: Python :: 3.14",
        "Topic :: Scientific/Engineering :: Quantum Computing",
    ],
    keywords=[
        "algorithms",
        "color code",
        "fault-tolerant quantum computing",
        "möbius decoder",
        "quantum computing",
        "quantum error correction",
        "quantum",
    ],
)