toploc/setup.py at main · PrimeIntellect-ai/toploc · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
from setuptools import setup
import os
import platform
from torch.utils.cpp_extension import BuildExtension, CppExtension

CSRC_DIR = os.path.join("toploc", "C", "csrc")

# Define mac-specific compiler and linker flags
if os.environ.get("DEBUG"):
    extra_compile_args = ["-DDEBUG", "-O0"]
else:
    extra_compile_args = ["-O3"]

extra_link_args: list[str] = []

# Add macOS specific flags
if platform.system() == "Darwin":
    # Enable support for both Intel and Apple Silicon
    extra_compile_args.extend(["-arch", "x86_64", "-arch", "arm64"])
    extra_link_args.extend(["-arch", "x86_64", "-arch", "arm64"])

    # Add minimum deployment target for macOS
    extra_compile_args.append("-mmacosx-version-min=10.13")
    extra_link_args.append("-mmacosx-version-min=10.13")

extensions = [
    CppExtension(
        name="toploc.C.csrc.ndd",
        sources=[os.path.join(CSRC_DIR, "ndd.cpp")],
        extra_compile_args=extra_compile_args,
        extra_link_args=extra_link_args,
    ),
    CppExtension(
        name="toploc.C.csrc.utils",
        sources=[os.path.join(CSRC_DIR, "utils.cpp")],
        extra_compile_args=extra_compile_args,
        extra_link_args=extra_link_args,
    ),
    CppExtension(
        name="toploc.C.csrc.poly",
        sources=[os.path.join(CSRC_DIR, "poly.cpp")],
        extra_compile_args=extra_compile_args,
        extra_link_args=extra_link_args,
    ),
]

setup(
    name="toploc",
    ext_modules=extensions,
    packages=["toploc", "toploc.C.csrc"],
    package_data={
        "toploc.C.csrc": ["*.pyi"],  # Include .pyi files
    },
    cmdclass={"build_ext": BuildExtension},
)