1import sys, os, platform
2
3from distutils.util  import split_quoted
4from distutils.spawn import find_executable
5from distutils import log as dulog
6
7try:
8    from collections import OrderedDict
9except ImportError:
10    OrderedDict = dict
11
12try:
13    from configparser import ConfigParser
14    from configparser import Error as ConfigParserError
15except ImportError:
16    from ConfigParser import ConfigParser
17    from ConfigParser import Error as ConfigParserError
18
19class Config(object):
20
21    def __init__(self, logger=None):
22        self.log = logger or dulog
23        self.section  = None
24        self.filename = None
25        self.compiler_info = OrderedDict((
26                ('mpicc'  , None),
27                ('mpicxx' , None),
28                ('mpifort', None),
29                ('mpif90' , None),
30                ('mpif77' , None),
31                ('mpild'  , None),
32                ))
33        self.library_info = OrderedDict((
34            ('define_macros'        , []),
35            ('undef_macros'         , []),
36            ('include_dirs'         , []),
37
38            ('libraries'            , []),
39            ('library_dirs'         , []),
40            ('runtime_library_dirs' , []),
41
42            ('extra_compile_args'   , []),
43            ('extra_link_args'      , []),
44            ('extra_objects'        , []),
45            ))
46
47    def __bool__(self):
48        for v in self.compiler_info.values():
49            if v:
50                return True
51        for v in self.library_info.values():
52            if v:
53                return True
54        return False
55
56    __nonzero__ = __bool__
57
58    def get(self, k, d=None):
59        if k in self.compiler_info:
60            return self.compiler_info[k]
61        if k in self.library_info:
62            return self.library_info[k]
63        return d
64
65    def info(self, log=None):
66        if log is None: log = self.log
67        mpicc   = self.compiler_info.get('mpicc')
68        mpicxx  = self.compiler_info.get('mpicxx')
69        mpifort = self.compiler_info.get('mpifort')
70        mpif90  = self.compiler_info.get('mpif90')
71        mpif77  = self.compiler_info.get('mpif77')
72        mpild   = self.compiler_info.get('mpild')
73        if mpicc:
74            log.info("MPI C compiler:    %s", mpicc)
75        if mpicxx:
76            log.info("MPI C++ compiler:  %s", mpicxx)
77        if mpifort:
78            log.info("MPI F compiler:    %s", mpifort)
79        if mpif90:
80            log.info("MPI F90 compiler:  %s", mpif90)
81        if mpif77:
82            log.info("MPI F77 compiler:  %s", mpif77)
83        if mpild:
84            log.info("MPI linker:        %s", mpild)
85
86    def update(self, config, **more):
87        if hasattr(config, 'keys'):
88            config = config.items()
89        for option, value in config:
90            if option in self.compiler_info:
91                self.compiler_info[option] = value
92            if option in self.library_info:
93                self.library_info[option] = value
94        if more:
95            self.update(more)
96
97    def setup(self, options, environ=None):
98        if environ is None: environ = os.environ
99        self.setup_library_info(options, environ)
100        self.setup_compiler_info(options, environ)
101
102    def setup_library_info(self, options, environ):
103        filename = section = None
104        mpiopt = getattr(options, 'mpi', None)
105        mpiopt = environ.get('MPICFG', mpiopt)
106        if mpiopt:
107            if ',' in mpiopt:
108                section, filename = mpiopt.split(',', 1)
109            elif os.path.isfile(mpiopt):
110                filename = mpiopt
111            else:
112                section = mpiopt
113        if not filename: filename = "mpi.cfg"
114        if not section:  section  = "mpi"
115
116        mach = platform.machine()
117        arch = platform.architecture()[0]
118        plat = sys.platform
119        osnm = os.name
120        if   'linux' == plat[:5]: plat = 'linux'
121        elif 'sunos' == plat[:5]: plat = 'solaris'
122        elif 'win'   == plat[:3]: plat = 'windows'
123        suffixes = []
124        suffixes.append(plat+'-'+mach)
125        suffixes.append(plat+'-'+arch)
126        suffixes.append(plat)
127        suffixes.append(osnm+'-'+mach)
128        suffixes.append(osnm+'-'+arch)
129        suffixes.append(osnm)
130        suffixes.append(mach)
131        suffixes.append(arch)
132        sections  = [section+"-"+s for s in suffixes]
133        sections += [section]
134        self.load(filename, sections)
135        if not self:
136            if os.name == 'posix':
137                self._setup_posix()
138            if sys.platform == 'win32':
139                self._setup_windows()
140
141    def _setup_posix(self):
142        pass
143
144    def _setup_windows(self):
145        if 'I_MPI_ROOT' in os.environ:
146            self._setup_windows_intelmpi()
147        else:
148            self._setup_windows_msmpi()
149
150    def _setup_windows_intelmpi(self):
151        I_MPI_ROOT = os.environ.get('I_MPI_ROOT', '')
152        if not I_MPI_ROOT: return
153        if not os.path.isdir(I_MPI_ROOT): return
154        arch = platform.architecture()[0][:2]
155        archdir = {'32':'ia32', '64':'intel64'}[arch]
156        mpi_dir = os.path.join(I_MPI_ROOT, archdir)
157        if not os.path.isdir(mpi_dir): return
158        IMPI_INC = os.path.join(mpi_dir, 'include')
159        IMPI_LIB = os.path.join(mpi_dir, 'lib', 'release')
160        if not os.path.isdir(IMPI_INC): return
161        if not os.path.isdir(IMPI_LIB): return
162        self.library_info.update(
163            include_dirs=[IMPI_INC],
164            library_dirs=[IMPI_LIB],
165            libraries=['impi'])
166        self.section = 'intelmpi'
167        self.filename = [mpi_dir]
168        return True
169
170    def _setup_windows_msmpi(self):
171        # Microsoft MPI (v7, v6, v5, v4)
172        def msmpi_ver():
173            try:
174                try:
175                    import winreg
176                except ImportError:
177                    import _winreg as winreg
178                HKLM = winreg.HKEY_LOCAL_MACHINE
179                subkey = r"SOFTWARE\Microsoft\MPI"
180                with winreg.OpenKey(HKLM, subkey) as key:
181                    for i in range(winreg.QueryInfoKey(key)[1]):
182                        name, value, type = winreg.EnumValue(key, i)
183                        if name != "Version": continue
184                        major, minor = value.split('.')[:2]
185                        return (int(major), int(minor))
186            except: pass
187            return (1, 0)
188        def setup_msmpi(MSMPI_INC, MSMPI_LIB):
189            from os.path import join, isfile
190            ok = (MSMPI_INC and isfile(join(MSMPI_INC, 'mpi.h')) and
191                  MSMPI_LIB and isfile(join(MSMPI_LIB, 'msmpi.lib')))
192            if not ok: return False
193            major, minor = msmpi_ver()
194            MSMPI_VER = hex((major<<8)|(minor&0xFF))
195            MSMPI_INC = os.path.normpath(MSMPI_INC)
196            MSMPI_LIB = os.path.normpath(MSMPI_LIB)
197            self.library_info.update(
198                define_macros=[('MSMPI_VER', MSMPI_VER)],
199                include_dirs=[MSMPI_INC],
200                library_dirs=[MSMPI_LIB],
201                libraries=['msmpi'])
202            self.section = 'msmpi'
203            self.filename = [os.path.dirname(MSMPI_INC)]
204            return True
205        arch = platform.architecture()[0][:2]
206        # Look for Microsoft MPI in the environment
207        MSMPI_INC = os.environ.get('MSMPI_INC')
208        MSMPI_LIB = os.environ.get('MSMPI_LIB'+arch)
209        if setup_msmpi(MSMPI_INC, MSMPI_LIB): return
210        # Look for Microsoft MPI v7/v6/v5 in default install path
211        for ProgramFiles in ('ProgramFiles', 'ProgramFiles(x86)'):
212            ProgramFiles = os.environ.get(ProgramFiles, '')
213            archdir = {'32':'x86', '64':'x64'}[arch]
214            MSMPI_DIR = os.path.join(ProgramFiles, 'Microsoft SDKs', 'MPI')
215            MSMPI_INC = os.path.join(MSMPI_DIR, 'Include')
216            MSMPI_LIB = os.path.join(MSMPI_DIR, 'Lib', archdir)
217            if setup_msmpi(MSMPI_INC, MSMPI_LIB): return
218        # Look for Microsoft HPC Pack 2012 R2 in default install path
219        for ProgramFiles in ('ProgramFiles', 'ProgramFiles(x86)'):
220            ProgramFiles = os.environ.get(ProgramFiles, '')
221            archdir = {'32':'i386', '64':'amd64'}[arch]
222            MSMPI_DIR = os.path.join(ProgramFiles, 'Microsoft MPI')
223            MSMPI_INC = os.path.join(MSMPI_DIR, 'Inc')
224            MSMPI_LIB = os.path.join(MSMPI_DIR, 'Lib', archdir)
225            if setup_msmpi(MSMPI_INC, MSMPI_LIB): return
226
227        # Microsoft MPI (legacy) and others
228        from glob import glob
229        ProgramFiles = os.environ.get('ProgramFiles', '')
230        CCP_HOME = os.environ.get('CCP_HOME', '')
231        for (name, prefix, suffix) in (
232            ('msmpi',    CCP_HOME,     ''),
233            ('msmpi',    ProgramFiles, 'Microsoft HPC Pack 2012 R2'),
234            ('msmpi',    ProgramFiles, 'Microsoft HPC Pack 2012'),
235            ('msmpi',    ProgramFiles, 'Microsoft HPC Pack 2012 SDK'),
236            ('msmpi',    ProgramFiles, 'Microsoft HPC Pack 2008 R2'),
237            ('msmpi',    ProgramFiles, 'Microsoft HPC Pack 2008'),
238            ('msmpi',    ProgramFiles, 'Microsoft HPC Pack 2008 SDK'),
239            ):
240            mpi_dir = os.path.join(prefix, suffix)
241            if not mpi_dir or not os.path.isdir(mpi_dir): continue
242            define_macros = []
243            include_dir = os.path.join(mpi_dir, 'include')
244            library = 'mpi'
245            library_dir = os.path.join(mpi_dir, 'lib')
246            if name == 'msmpi':
247                include_dir = os.path.join(mpi_dir, 'inc')
248                library = 'msmpi'
249                arch = platform.architecture()[0]
250                if arch == '32bit':
251                    library_dir = os.path.join(library_dir, 'i386')
252                if arch == '64bit':
253                    library_dir = os.path.join(library_dir, 'amd64')
254                if not os.path.isdir(include_dir):
255                    include_dir = os.path.join(mpi_dir, 'include')
256            self.library_info.update(
257                define_macros=define_macros,
258                include_dirs=[include_dir],
259                libraries=[library],
260                library_dirs=[library_dir],
261                )
262            self.section = name
263            self.filename = [mpi_dir]
264            break
265
266
267    def setup_compiler_info(self, options, environ):
268        def find_exe(cmd, path=None):
269            if not cmd: return None
270            parts = split_quoted(cmd)
271            exe, args = parts[0], parts[1:]
272            if not os.path.isabs(exe) and path:
273                exe = os.path.basename(exe)
274            exe = find_executable(exe, path)
275            if not exe: return None
276            return ' '.join([exe]+args)
277        COMPILERS = (
278            ('mpicc',   ['mpicc',   'mpcc_r']),
279            ('mpicxx',  ['mpicxx',  'mpic++', 'mpiCC', 'mpCC_r']),
280            ('mpifort', ['mpifort', 'mpif90', 'mpif77', 'mpfort_r']),
281            ('mpif90',  ['mpif90',  'mpf90_r']),
282            ('mpif77',  ['mpif77',  'mpf77_r']),
283            ('mpild',   []),
284            )
285        #
286        compiler_info = {}
287        PATH = environ.get('PATH', '')
288        for name, _ in COMPILERS:
289            cmd = (environ.get(name.upper()) or
290                   getattr(options, name, None) or
291                   self.compiler_info.get(name) or
292                   None)
293            if cmd:
294                exe = find_exe(cmd, path=PATH)
295                if exe:
296                    path = os.path.dirname(exe)
297                    PATH = path + os.path.pathsep + PATH
298                    compiler_info[name] = exe
299                else:
300                    self.log.error("error: '%s' not found", cmd)
301        #
302        if not self and not compiler_info:
303            for name, candidates in COMPILERS:
304                for cmd in candidates:
305                    cmd = find_exe(cmd)
306                    if cmd:
307                        compiler_info[name] = cmd
308                        break
309        #
310        self.compiler_info.update(compiler_info)
311
312
313    def load(self, filename="mpi.cfg", section='mpi'):
314        if isinstance(filename, str):
315            filenames = filename.split(os.path.pathsep)
316        else:
317            filenames = list(filename)
318        if isinstance(section, str):
319            sections = section.split(',')
320        else:
321            sections = list(section)
322        #
323        try:
324            parser = ConfigParser(dict_type=OrderedDict)
325        except TypeError:
326            parser = ConfigParser()
327        try:
328            read_ok = parser.read(filenames)
329        except ConfigParserError:
330            self.log.error(
331                "error: parsing configuration file/s '%s'",
332                os.path.pathsep.join(filenames))
333            return None
334        for section in sections:
335            if parser.has_section(section):
336                break
337            section = None
338        if not section:
339            self.log.error(
340                "error: section/s '%s' not found in file/s '%s'",
341                ','.join(sections), os.path.pathsep.join(filenames))
342            return None
343        parser_items = list(parser.items(section, vars=None))
344        #
345        compiler_info = type(self.compiler_info)()
346        for option, value in parser_items:
347            if option in self.compiler_info:
348                compiler_info[option] = value
349        #
350        pathsep = os.path.pathsep
351        expanduser = os.path.expanduser
352        expandvars = os.path.expandvars
353        library_info = type(self.library_info)()
354        for k, v in parser_items:
355            if k in ('define_macros',
356                     'undef_macros',
357                     ):
358                macros = [e.strip() for e in v.split(',')]
359                if k == 'define_macros':
360                    for i, m in enumerate(macros):
361                        try: # -DFOO=bar
362                            idx = m.index('=')
363                            macro = (m[:idx], m[idx+1:] or None)
364                        except ValueError: # -DFOO
365                            macro = (m, None)
366                        macros[i] = macro
367                library_info[k] = macros
368            elif k in ('include_dirs',
369                       'library_dirs',
370                       'runtime_dirs',
371                       'runtime_library_dirs',
372                       ):
373                if k == 'runtime_dirs': k = 'runtime_library_dirs'
374                pathlist = [p.strip() for p in v.split(pathsep)]
375                library_info[k] = [expanduser(expandvars(p))
376                                   for p in pathlist if p]
377            elif k == 'libraries':
378                library_info[k] = [e.strip() for e in split_quoted(v)]
379            elif k in ('extra_compile_args',
380                       'extra_link_args',
381                       ):
382                library_info[k] = split_quoted(v)
383            elif k == 'extra_objects':
384                library_info[k] = [expanduser(expandvars(e))
385                                   for e in split_quoted(v)]
386            elif hasattr(self, k):
387                library_info[k] = v.strip()
388            else:
389                pass
390        #
391        self.section = section
392        self.filename = read_ok
393        self.compiler_info.update(compiler_info)
394        self.library_info.update(library_info)
395        return compiler_info, library_info, section, read_ok
396
397    def dump(self, filename=None, section='mpi'):
398        # prepare configuration values
399        compiler_info   = self.compiler_info.copy()
400        library_info = self.library_info.copy()
401        for k in library_info:
402            if k in ('define_macros',
403                     'undef_macros',
404                     ):
405                macros = library_info[k]
406                if k == 'define_macros':
407                    for i, (m, v) in enumerate(macros):
408                        if v is None:
409                            macros[i] = m
410                        else:
411                            macros[i] = '%s=%s' % (m, v)
412                library_info[k] = ','.join(macros)
413            elif k in ('include_dirs',
414                       'library_dirs',
415                       'runtime_library_dirs',
416                       ):
417                library_info[k] = os.path.pathsep.join(library_info[k])
418            elif isinstance(library_info[k], list):
419                library_info[k] = ' '.join(library_info[k])
420        # fill configuration parser
421        try:
422            parser = ConfigParser(dict_type=OrderedDict)
423        except TypeError:
424            parser = ConfigParser()
425        parser.add_section(section)
426        for option, value in compiler_info.items():
427            if not value: continue
428            parser.set(section, option, value)
429        for option, value in library_info.items():
430            if not value: continue
431            parser.set(section, option, value)
432        # save configuration file
433        if filename is None:
434            parser.write(sys.stdout)
435        elif hasattr(filename, 'write'):
436            parser.write(filename)
437        elif isinstance(filename, str):
438            with open(filename, 'w') as f:
439                parser.write(f)
440        return parser
441
442
443if __name__ == '__main__':
444    import optparse
445    parser = optparse.OptionParser()
446    parser.add_option("--mpi",     type="string")
447    parser.add_option("--mpicc",   type="string")
448    parser.add_option("--mpicxx",  type="string")
449    parser.add_option("--mpifort", type="string")
450    parser.add_option("--mpif90",  type="string")
451    parser.add_option("--mpif77",  type="string")
452    parser.add_option("--mpild",   type="string")
453    (opts, args) = parser.parse_args()
454
455    log = dulog.Log(dulog.INFO)
456    cfg = Config(log)
457    cfg.setup(opts)
458    cfg.dump()
459