Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""
2Build a c-extension module on-the-fly in tests.
3See build_and_import_extensions for usage hints
5"""
7import os
8import pathlib
9import sys
10import sysconfig
12__all__ = ['build_and_import_extension', 'compile_extension_module']
15def build_and_import_extension(
16 modname, functions, *, prologue="", build_dir=None,
17 include_dirs=[], more_init=""):
18 """
19 Build and imports a c-extension module `modname` from a list of function
20 fragments `functions`.
23 Parameters
24 ----------
25 functions : list of fragments
26 Each fragment is a sequence of func_name, calling convention, snippet.
27 prologue : string
28 Code to precede the rest, usually extra ``#include`` or ``#define``
29 macros.
30 build_dir : pathlib.Path
31 Where to build the module, usually a temporary directory
32 include_dirs : list
33 Extra directories to find include files when compiling
34 more_init : string
35 Code to appear in the module PyMODINIT_FUNC
37 Returns
38 -------
39 out: module
40 The module will have been loaded and is ready for use
42 Examples
43 --------
44 >>> functions = [("test_bytes", "METH_O", \"\"\"
45 if ( !PyBytesCheck(args)) {
46 Py_RETURN_FALSE;
47 }
48 Py_RETURN_TRUE;
49 \"\"\")]
50 >>> mod = build_and_import_extension("testme", functions)
51 >>> assert not mod.test_bytes(u'abc')
52 >>> assert mod.test_bytes(b'abc')
53 """
54 from distutils.errors import CompileError
56 body = prologue + _make_methods(functions, modname)
57 init = """PyObject *mod = PyModule_Create(&moduledef);
58 """
59 if not build_dir:
60 build_dir = pathlib.Path('.')
61 if more_init:
62 init += """#define INITERROR return NULL
63 """
64 init += more_init
65 init += "\nreturn mod;"
66 source_string = _make_source(modname, init, body)
67 try:
68 mod_so = compile_extension_module(
69 modname, build_dir, include_dirs, source_string)
70 except CompileError as e:
71 # shorten the exception chain
72 raise RuntimeError(f"could not compile in {build_dir}:") from e
73 import importlib.util
74 spec = importlib.util.spec_from_file_location(modname, mod_so)
75 foo = importlib.util.module_from_spec(spec)
76 spec.loader.exec_module(foo)
77 return foo
80def compile_extension_module(
81 name, builddir, include_dirs,
82 source_string, libraries=[], library_dirs=[]):
83 """
84 Build an extension module and return the filename of the resulting
85 native code file.
87 Parameters
88 ----------
89 name : string
90 name of the module, possibly including dots if it is a module inside a
91 package.
92 builddir : pathlib.Path
93 Where to build the module, usually a temporary directory
94 include_dirs : list
95 Extra directories to find include files when compiling
96 libraries : list
97 Libraries to link into the extension module
98 library_dirs: list
99 Where to find the libraries, ``-L`` passed to the linker
100 """
101 modname = name.split('.')[-1]
102 dirname = builddir / name
103 dirname.mkdir(exist_ok=True)
104 cfile = _convert_str_to_file(source_string, dirname)
105 include_dirs = include_dirs + [sysconfig.get_config_var('INCLUDEPY')]
107 return _c_compile(
108 cfile, outputfilename=dirname / modname,
109 include_dirs=include_dirs, libraries=[], library_dirs=[],
110 )
113def _convert_str_to_file(source, dirname):
114 """Helper function to create a file ``source.c`` in `dirname` that contains
115 the string in `source`. Returns the file name
116 """
117 filename = dirname / 'source.c'
118 with filename.open('w') as f:
119 f.write(str(source))
120 return filename
123def _make_methods(functions, modname):
124 """ Turns the name, signature, code in functions into complete functions
125 and lists them in a methods_table. Then turns the methods_table into a
126 ``PyMethodDef`` structure and returns the resulting code fragment ready
127 for compilation
128 """
129 methods_table = []
130 codes = []
131 for funcname, flags, code in functions:
132 cfuncname = "%s_%s" % (modname, funcname)
133 if 'METH_KEYWORDS' in flags:
134 signature = '(PyObject *self, PyObject *args, PyObject *kwargs)'
135 else:
136 signature = '(PyObject *self, PyObject *args)'
137 methods_table.append(
138 "{\"%s\", (PyCFunction)%s, %s}," % (funcname, cfuncname, flags))
139 func_code = """
140 static PyObject* {cfuncname}{signature}
141 {{
142 {code}
143 }}
144 """.format(cfuncname=cfuncname, signature=signature, code=code)
145 codes.append(func_code)
147 body = "\n".join(codes) + """
148 static PyMethodDef methods[] = {
149 %(methods)s
150 { NULL }
151 };
152 static struct PyModuleDef moduledef = {
153 PyModuleDef_HEAD_INIT,
154 "%(modname)s", /* m_name */
155 NULL, /* m_doc */
156 -1, /* m_size */
157 methods, /* m_methods */
158 };
159 """ % dict(methods='\n'.join(methods_table), modname=modname)
160 return body
163def _make_source(name, init, body):
164 """ Combines the code fragments into source code ready to be compiled
165 """
166 code = """
167 #include <Python.h>
169 %(body)s
171 PyMODINIT_FUNC
172 PyInit_%(name)s(void) {
173 %(init)s
174 }
175 """ % dict(
176 name=name, init=init, body=body,
177 )
178 return code
181def _c_compile(cfile, outputfilename, include_dirs=[], libraries=[],
182 library_dirs=[]):
183 if sys.platform == 'win32':
184 compile_extra = ["/we4013"]
185 link_extra = ["/LIBPATH:" + os.path.join(sys.base_prefix, 'libs')]
186 elif sys.platform.startswith('linux'):
187 compile_extra = [
188 "-O0", "-g", "-Werror=implicit-function-declaration", "-fPIC"]
189 link_extra = None
190 else:
191 compile_extra = link_extra = None
192 pass
193 if sys.platform == 'win32':
194 link_extra = link_extra + ['/DEBUG'] # generate .pdb file
195 if sys.platform == 'darwin':
196 # support Fink & Darwinports
197 for s in ('/sw/', '/opt/local/'):
198 if (s + 'include' not in include_dirs
199 and os.path.exists(s + 'include')):
200 include_dirs.append(s + 'include')
201 if s + 'lib' not in library_dirs and os.path.exists(s + 'lib'):
202 library_dirs.append(s + 'lib')
204 outputfilename = outputfilename.with_suffix(get_so_suffix())
205 saved_environ = os.environ.copy()
206 try:
207 build(
208 cfile, outputfilename,
209 compile_extra, link_extra,
210 include_dirs, libraries, library_dirs)
211 finally:
212 # workaround for a distutils bugs where some env vars can
213 # become longer and longer every time it is used
214 for key, value in saved_environ.items():
215 if os.environ.get(key) != value:
216 os.environ[key] = value
217 return outputfilename
220def build(cfile, outputfilename, compile_extra, link_extra,
221 include_dirs, libraries, library_dirs):
222 "cd into the directory where the cfile is, use distutils to build"
223 from numpy.distutils.ccompiler import new_compiler
225 compiler = new_compiler(force=1, verbose=2)
226 compiler.customize('')
227 objects = []
229 old = os.getcwd()
230 os.chdir(cfile.parent)
231 try:
232 res = compiler.compile(
233 [str(cfile.name)],
234 include_dirs=include_dirs,
235 extra_preargs=compile_extra
236 )
237 objects += [str(cfile.parent / r) for r in res]
238 finally:
239 os.chdir(old)
241 compiler.link_shared_object(
242 objects, str(outputfilename),
243 libraries=libraries,
244 extra_preargs=link_extra,
245 library_dirs=library_dirs)
248def get_so_suffix():
249 ret = sysconfig.get_config_var('EXT_SUFFIX')
250 assert ret
251 return ret