|
23 | 23 | import platform |
24 | 24 | import shutil |
25 | 25 |
|
26 | | -from setuptools import setup |
| 26 | +from setuptools import setup, Extension |
| 27 | +from setuptools.command.build_ext import build_ext as _build_ext |
27 | 28 |
|
28 | 29 | # make the jecq python package dir |
29 | 30 | shutil.rmtree("jecq", ignore_errors=True) |
|
32 | 33 | shutil.copyfile("loader.py", "jecq/loader.py") |
33 | 34 | shutil.copyfile("class_wrappers.py", "jecq/class_wrappers.py") |
34 | 35 |
|
35 | | -ext = ".pyd" if platform.system() == "Windows" else ".so" |
| 36 | +is_windows = platform.system() == "Windows" |
36 | 37 |
|
| 38 | +ext = ".pyd" if is_windows else ".so" |
37 | 39 | build_type = os.environ.get("JECQ_BUILD_TYPE", "RelWithDebInfo") |
38 | | -prefix = f"{build_type}/" * (platform.system() == "Windows") |
| 40 | +prefix = f"{build_type}/" * is_windows |
39 | 41 |
|
40 | 42 | swigjecq_generic_lib = f"{prefix}_swigjecq{ext}" |
41 | 43 | swigjecq_avx2_lib = f"{prefix}_swigjecq_avx2{ext}" |
42 | 44 | swigjecq_avx512_lib = f"{prefix}_swigjecq_avx512{ext}" |
43 | 45 | swigjecq_avx512_spr_lib = f"{prefix}_swigjecq_avx512_spr{ext}" |
44 | | -callbacks_lib = f"{prefix}libfaiss_python_callbacks{ext}" |
45 | 46 | swigjecq_sve_lib = f"{prefix}_swigjecq_sve{ext}" |
46 | 47 |
|
47 | 48 | found_swigjecq_generic = os.path.exists(swigjecq_generic_lib) |
48 | 49 | found_swigjecq_avx2 = os.path.exists(swigjecq_avx2_lib) |
49 | 50 | found_swigjecq_avx512 = os.path.exists(swigjecq_avx512_lib) |
50 | 51 | found_swigjecq_avx512_spr = os.path.exists(swigjecq_avx512_spr_lib) |
51 | | -found_callbacks = os.path.exists(callbacks_lib) |
52 | 52 | found_swigjecq_sve = os.path.exists(swigjecq_sve_lib) |
53 | 53 |
|
54 | 54 | assert ( |
|
64 | 64 | f"Jecq may not be compiled yet." |
65 | 65 | ) |
66 | 66 |
|
| 67 | +libs = [] |
| 68 | + |
67 | 69 | if found_swigjecq_generic: |
68 | 70 | print(f"Copying {swigjecq_generic_lib}") |
69 | 71 | shutil.copyfile("swigjecq.py", "jecq/swigjecq.py") |
70 | 72 | shutil.copyfile(swigjecq_generic_lib, f"jecq/_swigjecq{ext}") |
| 73 | + libs.append("_swigjecq") |
71 | 74 |
|
72 | 75 | if found_swigjecq_avx2: |
73 | 76 | print(f"Copying {swigjecq_avx2_lib}") |
74 | 77 | shutil.copyfile("swigjecq_avx2.py", "jecq/swigjecq_avx2.py") |
75 | 78 | shutil.copyfile(swigjecq_avx2_lib, f"jecq/_swigjecq_avx2{ext}") |
| 79 | + libs.append("_swigjecq_avx2") |
76 | 80 |
|
77 | 81 | if found_swigjecq_avx512: |
78 | 82 | print(f"Copying {swigjecq_avx512_lib}") |
79 | 83 | shutil.copyfile("swigjecq_avx512.py", "jecq/swigjecq_avx512.py") |
80 | 84 | shutil.copyfile(swigjecq_avx512_lib, f"jecq/_swigjecq_avx512{ext}") |
| 85 | + libs.append("_swigjecq_avx512") |
| 86 | + |
81 | 87 |
|
82 | 88 | if found_swigjecq_avx512_spr: |
83 | 89 | print(f"Copying {swigjecq_avx512_spr_lib}") |
84 | 90 | shutil.copyfile("swigjecq_avx512_spr.py", "jecq/swigjecq_avx512_spr.py") |
85 | 91 | shutil.copyfile(swigjecq_avx512_spr_lib, f"jecq/_swigjecq_avx512_spr{ext}") |
| 92 | + libs.append("_swigjecq_avx512_spr") |
86 | 93 |
|
87 | | -if found_callbacks: |
88 | | - print(f"Copying {callbacks_lib}") |
89 | | - shutil.copyfile(callbacks_lib, f"jecq/{callbacks_lib}") |
90 | 94 |
|
91 | 95 | if found_swigjecq_sve: |
92 | 96 | print(f"Copying {swigjecq_sve_lib}") |
93 | 97 | shutil.copyfile("swigjecq_sve.py", "jecq/swigjecq_sve.py") |
94 | 98 | shutil.copyfile(swigjecq_sve_lib, f"jecq/_swigjecq_sve{ext}") |
| 99 | + libs.append("_swigjecq_sve") |
| 100 | + |
| 101 | +ext_modules = [Extension(f"jecq.{mod}", sources=[]) for mod in libs] |
| 102 | + |
| 103 | + |
| 104 | +# 2) CopyBuildExt just copies the .so into the build directory without compiling |
| 105 | +class CopyBuildExt(_build_ext): |
| 106 | + def run(self): |
| 107 | + # skip the normal compiler run |
| 108 | + for extension in self.extensions: |
| 109 | + self.build_extension(extension) |
| 110 | + |
| 111 | + def build_extension(self, extension): |
| 112 | + name = extension.name.split(".", 1)[1] # e.g. "_swigjecq" |
| 113 | + src = f"{prefix}{name}{ext}" |
| 114 | + if not os.path.exists(src): |
| 115 | + raise FileNotFoundError(f"Binary output '{src}' not found") |
| 116 | + dst = self.get_ext_fullpath(extension.name) |
| 117 | + self.mkpath(os.path.dirname(dst)) |
| 118 | + shutil.copyfile(src, dst) |
| 119 | + print(f"Copied {src} to {dst}") |
| 120 | + |
95 | 121 |
|
96 | 122 | long_description = """ |
97 | 123 | Jecq is a library for efficient similarity search based on the Faiss library. |
98 | 124 | """ |
| 125 | + |
| 126 | +try: |
| 127 | + from wheel.bdist_wheel import bdist_wheel as _bdist_wheel |
| 128 | + |
| 129 | + """Custom bdist_wheel to ensure that the root is not pure Python.""" |
| 130 | + |
| 131 | + class bdist_wheel(_bdist_wheel): |
| 132 | + def finalize_options(self): |
| 133 | + _bdist_wheel.finalize_options(self) |
| 134 | + self.root_is_pure = False |
| 135 | + |
| 136 | +except ImportError: |
| 137 | + print("Not using custom bdist_wheel; wheel package not installed") |
| 138 | + |
99 | 139 | setup( |
100 | 140 | name="jecq", |
101 | 141 | version="0.0.1", |
102 | 142 | description="A Faiss-based library for efficient similarity search " |
103 | 143 | "and clustering of dense vectors", |
104 | 144 | long_description=long_description, |
105 | | - license="TODO", |
| 145 | + url="https://github.com/JaneaSystems/jecq", |
| 146 | + license="MIT", |
106 | 147 | keywords="search nearest neighbors", |
107 | 148 | install_requires=["numpy", "packaging", "faiss-cpu"], |
108 | 149 | packages=["jecq"], |
109 | 150 | package_data={ |
110 | | - "jecq": ["*.so", "*.pyd"], |
| 151 | + "jecq": ["*.pyd"] if is_windows else [], |
111 | 152 | }, |
112 | 153 | zip_safe=False, |
| 154 | + ext_modules=[] if is_windows else ext_modules, |
| 155 | + cmdclass={"bdist_wheel": bdist_wheel, "build_ext": CopyBuildExt}, |
113 | 156 | ) |
0 commit comments