1
- from setuptools import setup
2
- from setuptools .command .develop import develop
3
- from setuptools .command .install import install
1
+ import os
2
+ import re
3
+ import subprocess
4
+ import sys
5
+ from pathlib import Path
6
+
7
+ from setuptools import Extension , setup , find_packages
4
8
from setuptools .command .build_ext import build_ext
5
9
6
- class PostDevelopCommand (develop ):
7
- """Post-installation for development mode."""
8
- def run (self ):
9
- develop .run (self )
10
- print ("TensorArray developed." )
11
- # PUT YOUR POST-INSTALL SCRIPT HERE or CALL A FUNCTION
12
-
13
- class PostInstallCommand (install ):
14
- """Post-installation for installation mode."""
15
- def run (self ):
16
- install .run (self )
17
- print ("TensorArray installed." )
18
- # PUT YOUR POST-INSTALL SCRIPT HERE or CALL A FUNCTION
19
-
20
- setup (
21
- cmdclass = {
22
- 'develop' : PostDevelopCommand ,
23
- 'install' : PostInstallCommand ,
24
- },
25
- )
10
+ # Convert distutils Windows platform specifiers to CMake -A arguments
11
+ PLAT_TO_CMAKE = {
12
+ "win32" : "Win32" ,
13
+ "win-amd64" : "x64" ,
14
+ "win-arm32" : "ARM" ,
15
+ "win-arm64" : "ARM64" ,
16
+ }
17
+
18
+ cwd = os .path .dirname (os .path .abspath (__file__ ))
19
+
20
+ # A CMakeExtension needs a sourcedir instead of a file list.
21
+ # The name must be the _single_ output extension from the CMake build.
22
+ # If you need multiple extensions, see scikit-build.
23
+ class CMakeExtension (Extension ):
24
+ def __init__ (self , name : str , sourcedir : str = "" ) -> None :
25
+ super ().__init__ (name , sources = [])
26
+ self .sourcedir = os .fspath (Path (sourcedir ).resolve ())
27
+
28
+
29
+ class CMakeBuild (build_ext ):
30
+ def build_extension (self , ext : CMakeExtension ) -> None :
31
+ # Must be in this form due to bug in .resolve() only fixed in Python 3.10+
32
+ ext_fullpath = Path .cwd () / self .get_ext_fullpath (ext .name )
33
+ extdir = ext_fullpath .parent .resolve ()
34
+
35
+ # Using this requires trailing slash for auto-detection & inclusion of
36
+ # auxiliary "native" libs
37
+
38
+ debug = int (os .environ .get ("DEBUG" , 0 )) if self .debug is None else self .debug
39
+ cfg = "Debug" if debug else "Release"
40
+
41
+ # CMake lets you override the generator - we need to check this.
42
+ # Can be set with Conda-Build, for example.
43
+ cmake_generator = os .environ .get ("CMAKE_GENERATOR" , "" )
44
+
45
+ # Set Python_EXECUTABLE instead if you use PYBIND11_FINDPYTHON
46
+ # EXAMPLE_VERSION_INFO shows you how to pass a value into the C++ code
47
+ # from Python.
48
+ cmake_args = [
49
+ f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={ extdir } { os .sep } " ,
50
+ f"-DPYTHON_EXECUTABLE={ sys .executable } " ,
51
+ f"-DCMAKE_BUILD_TYPE={ cfg } " , # not used on MSVC, but no harm
52
+ ]
53
+ build_args = []
54
+ # Adding CMake arguments set as environment variable
55
+ # (needed e.g. to build for ARM OSx on conda-forge)
56
+ if "CMAKE_ARGS" in os .environ :
57
+ cmake_args += [item for item in os .environ ["CMAKE_ARGS" ].split (" " ) if item ]
58
+
59
+ # In this example, we pass in the version to C++. You might not need to.
60
+ cmake_args += [f"-DEXAMPLE_VERSION_INFO={ self .distribution .get_version ()} " ]
61
+
62
+ if self .compiler .compiler_type != "msvc" :
63
+ # Using Ninja-build since it a) is available as a wheel and b)
64
+ # multithreads automatically. MSVC would require all variables be
65
+ # exported for Ninja to pick it up, which is a little tricky to do.
66
+ # Users can override the generator with CMAKE_GENERATOR in CMake
67
+ # 3.15+.
68
+ if not cmake_generator or cmake_generator == "Ninja" :
69
+ try :
70
+ import ninja
71
+
72
+ ninja_executable_path = Path (ninja .BIN_DIR ) / "ninja"
73
+ cmake_args += [
74
+ "-GNinja" ,
75
+ f"-DCMAKE_MAKE_PROGRAM:FILEPATH={ ninja_executable_path } " ,
76
+ ]
77
+ except ImportError :
78
+ pass
79
+
80
+ else :
81
+ # Single config generators are handled "normally"
82
+ single_config = any (x in cmake_generator for x in {"NMake" , "Ninja" })
83
+
84
+ # CMake allows an arch-in-generator style for backward compatibility
85
+ contains_arch = any (x in cmake_generator for x in {"ARM" , "Win64" })
86
+
87
+ # Specify the arch if using MSVC generator, but only if it doesn't
88
+ # contain a backward-compatibility arch spec already in the
89
+ # generator name.
90
+ if not single_config and not contains_arch :
91
+ cmake_args += ["-A" , PLAT_TO_CMAKE [self .plat_name ]]
92
+
93
+ # Multi-config generators have a different way to specify configs
94
+ if not single_config :
95
+ cmake_args += [
96
+ f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY_{ cfg .upper ()} ={ extdir } "
97
+ ]
98
+ build_args += ["--config" , cfg ]
99
+
100
+ if sys .platform .startswith ("darwin" ):
101
+ # Cross-compile support for macOS - respect ARCHFLAGS if set
102
+ archs = re .findall (r"-arch (\S+)" , os .environ .get ("ARCHFLAGS" , "" ))
103
+ if archs :
104
+ cmake_args += ["-DCMAKE_OSX_ARCHITECTURES={}" .format (";" .join (archs ))]
105
+
106
+ # Set CMAKE_BUILD_PARALLEL_LEVEL to control the parallel build level
107
+ # across all generators.
108
+ if "CMAKE_BUILD_PARALLEL_LEVEL" not in os .environ :
109
+ # self.parallel is a Python 3 only way to set parallel jobs by hand
110
+ # using -j in the build_ext call, not supported by pip or PyPA-build.
111
+ if hasattr (self , "parallel" ) and self .parallel :
112
+ # CMake 3.12+ only.
113
+ build_args += [f"-j{ self .parallel } " ]
114
+
115
+ build_temp = Path (self .build_temp ) / ext .name
116
+ if not build_temp .exists ():
117
+ build_temp .mkdir (parents = True )
118
+
119
+ subprocess .run (
120
+ ["cmake" , ext .sourcedir , * cmake_args ], cwd = build_temp , check = True
121
+ )
122
+ subprocess .run (
123
+ ["cmake" , "--build" , "." , * build_args ], cwd = build_temp , check = True
124
+ )
125
+
126
+ def main ():
127
+ with open (os .path .join (cwd , "README.md" ), encoding = "utf-8" ) as f :
128
+ long_description = f .read ()
129
+
130
+ setup (
131
+ name = "TensorArray" ,
132
+ version = "0.0.1a02" ,
133
+ description = "A machine learning package" ,
134
+ long_description = long_description ,
135
+ ext_modules = [CMakeExtension ('tensor_array' )],
136
+ authors = "TensorArray-Creators" ,
137
+ url = "https://github.com/Tensor-Array/Tensor-Array-Python" ,
138
+ packages = find_packages (exclude = ("tests.*" , "tests" ,)),
139
+ classifiers = [
140
+ "Development Status :: 2 - Pre-Alpha" ,
141
+
142
+ "Programming Language :: Python :: 3" ,
143
+ "Programming Language :: Python :: 3.9" ,
144
+ "Programming Language :: Python :: 3.10" ,
145
+ "Programming Language :: Python :: 3.11" ,
146
+ "Programming Language :: Python :: 3.12" ,
147
+ "Programming Language :: Python :: 3.13" ,
148
+
149
+ "License :: OSI Approved :: MIT License" ,
150
+
151
+ "Environment :: GPU :: NVIDIA CUDA :: 12" ,
152
+ ],
153
+ license = "MIT" ,
154
+ cmdclass = {
155
+ 'build_ext' : CMakeBuild
156
+ },
157
+ )
158
+
159
+ if __name__ == "__main__" :
160
+ main ()
0 commit comments