@@ -53,12 +53,13 @@ def get_version(pkg_path):
53
53
54
54
55
55
# Set some extra flags for compiling with coverage support
56
- extra_include_dirs = []
57
56
extra_compile_args = []
58
57
extra_link_args = []
59
58
extra_define_macros = []
60
- library_dirs = []
61
- libraries = []
59
+ device_include_dirs = []
60
+ device_library_dirs = []
61
+ device_libraries = []
62
+ device_define_macros = []
62
63
63
64
if os .getenv ("NANOARROW_PYTHON_COVERAGE" ) == "1" :
64
65
extra_compile_args .append ("--coverage" )
@@ -87,81 +88,57 @@ def get_version(pkg_path):
87
88
lib_dirs_err = ", " .join (f"'{ d } " for d in possible_libs )
88
89
raise ValueError (f"Can't find CUDA library directory. Checked { lib_dirs_err } " )
89
90
90
- extra_include_dirs .append (str (include_dir ))
91
- library_dirs .append (str (lib_dirs [0 ].parent ))
92
- libraries .append ("cuda" )
93
- extra_define_macros .append (("NANOARROW_DEVICE_WITH_CUDA" , 1 ))
91
+ device_include_dirs .append (str (include_dir ))
92
+ device_library_dirs .append (str (lib_dirs [0 ].parent ))
93
+ device_libraries .append ("cuda" )
94
+ device_define_macros .append (("NANOARROW_DEVICE_WITH_CUDA" , 1 ))
95
+
96
+
97
+ def nanoarrow_extension (
98
+ name , * , nanoarrow_c = False , nanoarrow_device = False , nanoarrow_ipc = False
99
+ ):
100
+ sources = ["src/" + name .replace ("." , "/" ) + ".pyx" ]
101
+ libraries = []
102
+ library_dirs = []
103
+ include_dirs = ["src/nanoarrow" , "vendor" ]
104
+ define_macros = list (extra_define_macros )
105
+
106
+ if nanoarrow_c :
107
+ sources .append ("vendor/nanoarrow.c" )
108
+
109
+ if nanoarrow_device :
110
+ sources .append ("vendor/nanoarrow_device.c" )
111
+ include_dirs .extend (device_include_dirs )
112
+ libraries .extend (device_libraries )
113
+ library_dirs .extend (device_library_dirs )
114
+ define_macros .extend (device_define_macros )
115
+
116
+ if nanoarrow_ipc :
117
+ sources .extend (["vendor/nanoarrow_ipc.c" , "vendor/flatcc.c" ])
118
+
119
+ return Extension (
120
+ name = name ,
121
+ include_dirs = include_dirs ,
122
+ language = "c" ,
123
+ sources = sources ,
124
+ extra_compile_args = extra_compile_args ,
125
+ extra_link_args = extra_link_args ,
126
+ define_macros = define_macros ,
127
+ library_dirs = library_dirs ,
128
+ libraries = libraries ,
129
+ )
94
130
95
131
96
132
setup (
97
133
ext_modules = [
98
- Extension (
99
- name = "nanoarrow._device" ,
100
- include_dirs = ["src/nanoarrow" , "vendor" ],
101
- language = "c" ,
102
- sources = [
103
- "src/nanoarrow/_device.pyx" ,
104
- "vendor/nanoarrow.c" ,
105
- "vendor/nanoarrow_device.c" ,
106
- ],
107
- extra_compile_args = extra_compile_args ,
108
- extra_link_args = extra_link_args ,
109
- define_macros = extra_define_macros ,
110
- library_dirs = library_dirs ,
111
- libraries = libraries ,
112
- ),
113
- Extension (
114
- name = "nanoarrow._types" ,
115
- include_dirs = ["src/nanoarrow" , "vendor" ],
116
- language = "c" ,
117
- sources = [
118
- "src/nanoarrow/_types.pyx" ,
119
- ],
120
- extra_compile_args = extra_compile_args ,
121
- extra_link_args = extra_link_args ,
122
- define_macros = extra_define_macros ,
123
- ),
124
- Extension (
125
- name = "nanoarrow._utils" ,
126
- include_dirs = extra_include_dirs + ["src/nanoarrow" , "vendor" ],
127
- language = "c" ,
128
- sources = [
129
- "src/nanoarrow/_utils.pyx" ,
130
- "vendor/nanoarrow.c" ,
131
- ],
132
- extra_compile_args = extra_compile_args ,
133
- extra_link_args = extra_link_args ,
134
- define_macros = extra_define_macros ,
135
- ),
136
- Extension (
137
- name = "nanoarrow._lib" ,
138
- include_dirs = extra_include_dirs + ["src/nanoarrow" , "vendor" ],
139
- language = "c" ,
140
- sources = [
141
- "src/nanoarrow/_lib.pyx" ,
142
- "vendor/nanoarrow.c" ,
143
- "vendor/nanoarrow_device.c" ,
144
- ],
145
- extra_compile_args = extra_compile_args ,
146
- extra_link_args = extra_link_args ,
147
- define_macros = extra_define_macros ,
148
- library_dirs = library_dirs ,
149
- libraries = libraries ,
150
- ),
151
- Extension (
152
- name = "nanoarrow._ipc_lib" ,
153
- include_dirs = extra_include_dirs + ["src/nanoarrow" , "vendor" ],
154
- language = "c" ,
155
- sources = [
156
- "src/nanoarrow/_ipc_lib.pyx" ,
157
- "vendor/nanoarrow.c" ,
158
- "vendor/nanoarrow_ipc.c" ,
159
- "vendor/flatcc.c" ,
160
- ],
161
- extra_compile_args = extra_compile_args ,
162
- extra_link_args = extra_link_args ,
163
- define_macros = extra_define_macros ,
134
+ nanoarrow_extension ("nanoarrow._types" ),
135
+ nanoarrow_extension ("nanoarrow._utils" , nanoarrow_c = True ),
136
+ nanoarrow_extension (
137
+ "nanoarrow._device" , nanoarrow_c = True , nanoarrow_device = True
164
138
),
139
+ nanoarrow_extension ("nanoarrow._buffer" , nanoarrow_c = True ),
140
+ nanoarrow_extension ("nanoarrow._lib" , nanoarrow_c = True , nanoarrow_device = True ),
141
+ nanoarrow_extension ("nanoarrow._ipc_lib" , nanoarrow_c = True , nanoarrow_ipc = True ),
165
142
],
166
143
version = version ,
167
144
)
0 commit comments