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
56
57
58
59
60
61
62
63
64
65
66
|
import os
import pathlib
import shlex
import shutil
import cffi
def _read_text(filename):
return pathlib.Path(filename).read_text(encoding="utf-8")
ffi = cffi.FFI()
ffi.set_source("libmpi", _read_text("libmpi.c.in"))
ffi.cdef(_read_text("libmpi.h"))
class mpicompiler:
from cffi import ffiplatform
def __init__(self, cc, ld=None):
self.cc = cc
self.ld = ld or cc
self.ffi_compile = self.ffiplatform.compile
def __enter__(self):
self.ffiplatform.compile = self.compile
def __exit__(self, *args):
self.ffiplatform.compile = self.ffi_compile
def configure(self, compiler):
def fix_command(command, cmd):
if not cmd:
return
cmd = shlex.split(cmd)
exe = shutil.which(cmd[0])
if not exe:
return
command[0] = exe
command += cmd[1:]
fix_command(compiler.compiler_so, self.cc)
fix_command(compiler.linker_so, self.ld)
def compile(self, *args, **kwargs):
from distutils.command import build_ext
customize_compiler_orig = build_ext.customize_compiler
def customize_compiler(compiler):
customize_compiler_orig(compiler)
self.configure(compiler)
build_ext.customize_compiler = customize_compiler
try:
return self.ffi_compile(*args, **kwargs)
finally:
build_ext.customize_compiler = customize_compiler_orig
if __name__ == "__main__":
cc = os.environ.get("MPICC", "mpicc")
ld = os.environ.get("MPILD")
with mpicompiler(cc, ld):
ffi.compile()
|
