Chris@87: """
Chris@87: Utility functions for
Chris@87: 
Chris@87: - building and importing modules on test time, using a temporary location
Chris@87: - detecting if compilers are present
Chris@87: 
Chris@87: """
Chris@87: from __future__ import division, absolute_import, print_function
Chris@87: 
Chris@87: import os
Chris@87: import sys
Chris@87: import subprocess
Chris@87: import tempfile
Chris@87: import shutil
Chris@87: import atexit
Chris@87: import textwrap
Chris@87: import re
Chris@87: import random
Chris@87: 
Chris@87: import nose
Chris@87: 
Chris@87: from numpy.compat import asbytes, asstr
Chris@87: import numpy.f2py
Chris@87: 
Chris@87: try:
Chris@87:     from hashlib import md5
Chris@87: except ImportError:
Chris@87:     from md5 import new as md5
Chris@87: 
Chris@87: #
Chris@87: # Maintaining a temporary module directory
Chris@87: #
Chris@87: 
Chris@87: _module_dir = None
Chris@87: 
Chris@87: def _cleanup():
Chris@87:     global _module_dir
Chris@87:     if _module_dir is not None:
Chris@87:         try:
Chris@87:             sys.path.remove(_module_dir)
Chris@87:         except ValueError:
Chris@87:             pass
Chris@87:         try:
Chris@87:             shutil.rmtree(_module_dir)
Chris@87:         except (IOError, OSError):
Chris@87:             pass
Chris@87:         _module_dir = None
Chris@87: 
Chris@87: def get_module_dir():
Chris@87:     global _module_dir
Chris@87:     if _module_dir is None:
Chris@87:         _module_dir = tempfile.mkdtemp()
Chris@87:         atexit.register(_cleanup)
Chris@87:         if _module_dir not in sys.path:
Chris@87:             sys.path.insert(0, _module_dir)
Chris@87:     return _module_dir
Chris@87: 
Chris@87: def get_temp_module_name():
Chris@87:     # Assume single-threaded, and the module dir usable only by this thread
Chris@87:     d = get_module_dir()
Chris@87:     for j in range(5403, 9999999):
Chris@87:         name = "_test_ext_module_%d" % j
Chris@87:         fn = os.path.join(d, name)
Chris@87:         if name not in sys.modules and not os.path.isfile(fn+'.py'):
Chris@87:             return name
Chris@87:     raise RuntimeError("Failed to create a temporary module name")
Chris@87: 
Chris@87: def _memoize(func):
Chris@87:     memo = {}
Chris@87:     def wrapper(*a, **kw):
Chris@87:         key = repr((a, kw))
Chris@87:         if key not in memo:
Chris@87:             try:
Chris@87:                 memo[key] = func(*a, **kw)
Chris@87:             except Exception as e:
Chris@87:                 memo[key] = e
Chris@87:                 raise
Chris@87:         ret = memo[key]
Chris@87:         if isinstance(ret, Exception):
Chris@87:             raise ret
Chris@87:         return ret
Chris@87:     wrapper.__name__ = func.__name__
Chris@87:     return wrapper
Chris@87: 
Chris@87: #
Chris@87: # Building modules
Chris@87: #
Chris@87: 
Chris@87: @_memoize
Chris@87: def build_module(source_files, options=[], skip=[], only=[], module_name=None):
Chris@87:     """
Chris@87:     Compile and import a f2py module, built from the given files.
Chris@87: 
Chris@87:     """
Chris@87: 
Chris@87:     code = ("import sys; sys.path = %s; import numpy.f2py as f2py2e; "
Chris@87:             "f2py2e.main()" % repr(sys.path))
Chris@87: 
Chris@87:     d = get_module_dir()
Chris@87: 
Chris@87:     # Copy files
Chris@87:     dst_sources = []
Chris@87:     for fn in source_files:
Chris@87:         if not os.path.isfile(fn):
Chris@87:             raise RuntimeError("%s is not a file" % fn)
Chris@87:         dst = os.path.join(d, os.path.basename(fn))
Chris@87:         shutil.copyfile(fn, dst)
Chris@87:         dst_sources.append(dst)
Chris@87: 
Chris@87:         fn = os.path.join(os.path.dirname(fn), '.f2py_f2cmap')
Chris@87:         if os.path.isfile(fn):
Chris@87:             dst = os.path.join(d, os.path.basename(fn))
Chris@87:             if not os.path.isfile(dst):
Chris@87:                 shutil.copyfile(fn, dst)
Chris@87: 
Chris@87:     # Prepare options
Chris@87:     if module_name is None:
Chris@87:         module_name = get_temp_module_name()
Chris@87:     f2py_opts = ['-c', '-m', module_name] + options + dst_sources
Chris@87:     if skip:
Chris@87:         f2py_opts += ['skip:'] + skip
Chris@87:     if only:
Chris@87:         f2py_opts += ['only:'] + only
Chris@87: 
Chris@87:     # Build
Chris@87:     cwd = os.getcwd()
Chris@87:     try:
Chris@87:         os.chdir(d)
Chris@87:         cmd = [sys.executable, '-c', code] + f2py_opts
Chris@87:         p = subprocess.Popen(cmd, stdout=subprocess.PIPE,
Chris@87:                              stderr=subprocess.STDOUT)
Chris@87:         out, err = p.communicate()
Chris@87:         if p.returncode != 0:
Chris@87:             raise RuntimeError("Running f2py failed: %s\n%s"
Chris@87:                                % (cmd[4:], asstr(out)))
Chris@87:     finally:
Chris@87:         os.chdir(cwd)
Chris@87: 
Chris@87:         # Partial cleanup
Chris@87:         for fn in dst_sources:
Chris@87:             os.unlink(fn)
Chris@87: 
Chris@87:     # Import
Chris@87:     __import__(module_name)
Chris@87:     return sys.modules[module_name]
Chris@87: 
Chris@87: @_memoize
Chris@87: def build_code(source_code, options=[], skip=[], only=[], suffix=None,
Chris@87:                module_name=None):
Chris@87:     """
Chris@87:     Compile and import Fortran code using f2py.
Chris@87: 
Chris@87:     """
Chris@87:     if suffix is None:
Chris@87:         suffix = '.f'
Chris@87: 
Chris@87:     fd, tmp_fn = tempfile.mkstemp(suffix=suffix)
Chris@87:     os.write(fd, asbytes(source_code))
Chris@87:     os.close(fd)
Chris@87: 
Chris@87:     try:
Chris@87:         return build_module([tmp_fn], options=options, skip=skip, only=only,
Chris@87:                             module_name=module_name)
Chris@87:     finally:
Chris@87:         os.unlink(tmp_fn)
Chris@87: 
Chris@87: #
Chris@87: # Check if compilers are available at all...
Chris@87: #
Chris@87: 
Chris@87: _compiler_status = None
Chris@87: def _get_compiler_status():
Chris@87:     global _compiler_status
Chris@87:     if _compiler_status is not None:
Chris@87:         return _compiler_status
Chris@87: 
Chris@87:     _compiler_status = (False, False, False)
Chris@87: 
Chris@87:     # XXX: this is really ugly. But I don't know how to invoke Distutils
Chris@87:     #      in a safer way...
Chris@87:     code = """
Chris@87: import os
Chris@87: import sys
Chris@87: sys.path = %(syspath)s
Chris@87: 
Chris@87: def configuration(parent_name='',top_path=None):
Chris@87:     global config
Chris@87:     from numpy.distutils.misc_util import Configuration
Chris@87:     config = Configuration('', parent_name, top_path)
Chris@87:     return config
Chris@87: 
Chris@87: from numpy.distutils.core import setup
Chris@87: setup(configuration=configuration)
Chris@87: 
Chris@87: config_cmd = config.get_config_cmd()
Chris@87: have_c = config_cmd.try_compile('void foo() {}')
Chris@87: print('COMPILERS:%%d,%%d,%%d' %% (have_c,
Chris@87:                                   config.have_f77c(),
Chris@87:                                   config.have_f90c()))
Chris@87: sys.exit(99)
Chris@87: """
Chris@87:     code = code % dict(syspath=repr(sys.path))
Chris@87: 
Chris@87:     fd, script = tempfile.mkstemp(suffix='.py')
Chris@87:     os.write(fd, asbytes(code))
Chris@87:     os.close(fd)
Chris@87: 
Chris@87:     try:
Chris@87:         cmd = [sys.executable, script, 'config']
Chris@87:         p = subprocess.Popen(cmd, stdout=subprocess.PIPE,
Chris@87:                              stderr=subprocess.STDOUT)
Chris@87:         out, err = p.communicate()
Chris@87:         m = re.search(asbytes(r'COMPILERS:(\d+),(\d+),(\d+)'), out)
Chris@87:         if m:
Chris@87:             _compiler_status = (bool(int(m.group(1))), bool(int(m.group(2))),
Chris@87:                                 bool(int(m.group(3))))
Chris@87:     finally:
Chris@87:         os.unlink(script)
Chris@87: 
Chris@87:     # Finished
Chris@87:     return _compiler_status
Chris@87: 
Chris@87: def has_c_compiler():
Chris@87:     return _get_compiler_status()[0]
Chris@87: 
Chris@87: def has_f77_compiler():
Chris@87:     return _get_compiler_status()[1]
Chris@87: 
Chris@87: def has_f90_compiler():
Chris@87:     return _get_compiler_status()[2]
Chris@87: 
Chris@87: #
Chris@87: # Building with distutils
Chris@87: #
Chris@87: 
Chris@87: @_memoize
Chris@87: def build_module_distutils(source_files, config_code, module_name, **kw):
Chris@87:     """
Chris@87:     Build a module via distutils and import it.
Chris@87: 
Chris@87:     """
Chris@87:     from numpy.distutils.misc_util import Configuration
Chris@87:     from numpy.distutils.core import setup
Chris@87: 
Chris@87:     d = get_module_dir()
Chris@87: 
Chris@87:     # Copy files
Chris@87:     dst_sources = []
Chris@87:     for fn in source_files:
Chris@87:         if not os.path.isfile(fn):
Chris@87:             raise RuntimeError("%s is not a file" % fn)
Chris@87:         dst = os.path.join(d, os.path.basename(fn))
Chris@87:         shutil.copyfile(fn, dst)
Chris@87:         dst_sources.append(dst)
Chris@87: 
Chris@87:     # Build script
Chris@87:     config_code = textwrap.dedent(config_code).replace("\n", "\n    ")
Chris@87: 
Chris@87:     code = """\
Chris@87: import os
Chris@87: import sys
Chris@87: sys.path = %(syspath)s
Chris@87: 
Chris@87: def configuration(parent_name='',top_path=None):
Chris@87:     from numpy.distutils.misc_util import Configuration
Chris@87:     config = Configuration('', parent_name, top_path)
Chris@87:     %(config_code)s
Chris@87:     return config
Chris@87: 
Chris@87: if __name__ == "__main__":
Chris@87:     from numpy.distutils.core import setup
Chris@87:     setup(configuration=configuration)
Chris@87: """ % dict(config_code=config_code, syspath = repr(sys.path))
Chris@87: 
Chris@87:     script = os.path.join(d, get_temp_module_name() + '.py')
Chris@87:     dst_sources.append(script)
Chris@87:     f = open(script, 'wb')
Chris@87:     f.write(asbytes(code))
Chris@87:     f.close()
Chris@87: 
Chris@87:     # Build
Chris@87:     cwd = os.getcwd()
Chris@87:     try:
Chris@87:         os.chdir(d)
Chris@87:         cmd = [sys.executable, script, 'build_ext', '-i']
Chris@87:         p = subprocess.Popen(cmd, stdout=subprocess.PIPE,
Chris@87:                              stderr=subprocess.STDOUT)
Chris@87:         out, err = p.communicate()
Chris@87:         if p.returncode != 0:
Chris@87:             raise RuntimeError("Running distutils build failed: %s\n%s"
Chris@87:                                % (cmd[4:], asstr(out)))
Chris@87:     finally:
Chris@87:         os.chdir(cwd)
Chris@87: 
Chris@87:         # Partial cleanup
Chris@87:         for fn in dst_sources:
Chris@87:             os.unlink(fn)
Chris@87: 
Chris@87:     # Import
Chris@87:     __import__(module_name)
Chris@87:     return sys.modules[module_name]
Chris@87: 
Chris@87: #
Chris@87: # Unittest convenience
Chris@87: #
Chris@87: 
Chris@87: class F2PyTest(object):
Chris@87:     code = None
Chris@87:     sources = None
Chris@87:     options = []
Chris@87:     skip = []
Chris@87:     only = []
Chris@87:     suffix = '.f'
Chris@87:     module = None
Chris@87:     module_name = None
Chris@87: 
Chris@87:     def setUp(self):
Chris@87:         if self.module is not None:
Chris@87:             return
Chris@87: 
Chris@87:         # Check compiler availability first
Chris@87:         if not has_c_compiler():
Chris@87:             raise nose.SkipTest("No C compiler available")
Chris@87: 
Chris@87:         codes = []
Chris@87:         if self.sources:
Chris@87:             codes.extend(self.sources)
Chris@87:         if self.code is not None:
Chris@87:             codes.append(self.suffix)
Chris@87: 
Chris@87:         needs_f77 = False
Chris@87:         needs_f90 = False
Chris@87:         for fn in codes:
Chris@87:             if fn.endswith('.f'):
Chris@87:                 needs_f77 = True
Chris@87:             elif fn.endswith('.f90'):
Chris@87:                 needs_f90 = True
Chris@87:         if needs_f77 and not has_f77_compiler():
Chris@87:             raise nose.SkipTest("No Fortran 77 compiler available")
Chris@87:         if needs_f90 and not has_f90_compiler():
Chris@87:             raise nose.SkipTest("No Fortran 90 compiler available")
Chris@87: 
Chris@87:         # Build the module
Chris@87:         if self.code is not None:
Chris@87:             self.module = build_code(self.code, options=self.options,
Chris@87:                                      skip=self.skip, only=self.only,
Chris@87:                                      suffix=self.suffix,
Chris@87:                                      module_name=self.module_name)
Chris@87: 
Chris@87:         if self.sources is not None:
Chris@87:             self.module = build_module(self.sources, options=self.options,
Chris@87:                                        skip=self.skip, only=self.only,
Chris@87:                                        module_name=self.module_name)