1import os
2import cffi
3
4ffi = cffi.FFI()
5with open("libmpi.c.in") as f:
6    ffi.set_source("libmpi", f.read())
7with open("libmpi.h") as f:
8    ffi.cdef(f.read())
9
10class mpicompiler(object):
11
12    from cffi import ffiplatform
13
14    def __init__(self, cc, ld=None):
15        self.cc = cc
16        self.ld = ld if ld else cc
17        self.ffi_compile = self.ffiplatform.compile
18
19    def __enter__(self):
20        self.ffiplatform.compile = self.compile
21
22    def __exit__(self, *args):
23        self.ffiplatform.compile = self.ffi_compile
24
25    def configure(self, compiler):
26        from distutils.util import split_quoted
27        from distutils.spawn import find_executable
28        def fix_command(command, cmd):
29            if not cmd: return
30            cmd = split_quoted(cmd)
31            exe = find_executable(cmd[0])
32            if not exe: return
33            command[0] = exe
34            command += cmd[1:]
35        fix_command(compiler.compiler_so, self.cc)
36        fix_command(compiler.linker_so, self.ld)
37
38    def compile(self, *args, **kargs):
39        from distutils.command import build_ext
40        customize_compiler_orig = build_ext.customize_compiler
41        def customize_compiler(compiler):
42            customize_compiler_orig(compiler)
43            self.configure(compiler)
44        build_ext.customize_compiler = customize_compiler
45        try:
46            return self.ffi_compile(*args, **kargs)
47        finally:
48            build_ext.customize_compiler = customize_compiler_orig
49
50if __name__ == '__main__':
51    cc = os.environ.get('MPICC', 'mpicc')
52    ld = os.environ.get('MPILD')
53    with mpicompiler(cc, ld):
54        ffi.compile()
55