File: setupext.py

package info (click to toggle)
yt 3.6.1-1
  • links: PTS, VCS
  • area: main
  • in suites: bookworm, bullseye, sid
  • size: 60,160 kB
  • sloc: python: 125,044; ansic: 9,049; sh: 626; makefile: 124; javascript: 91; csh: 36
file content (208 lines) | stat: -rw-r--r-- 6,271 bytes parent folder | download | duplicates (2)
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
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
import contextlib
import glob
import os
import shutil
import subprocess
import sys
import tempfile

from distutils import log
from distutils.ccompiler import new_compiler
from distutils.sysconfig import customize_compiler
from distutils.errors import CompileError, LinkError
from pkg_resources import resource_filename
from subprocess import Popen, PIPE


CCODE = """
#include <omp.h>
#include <stdio.h>
int main() {
  omp_set_num_threads(2);
  #pragma omp parallel
  printf("nthreads=%d\\n", omp_get_num_threads());
  return 0;
}
"""

@contextlib.contextmanager
def stdchannel_redirected(stdchannel, dest_filename):
    """
    A context manager to temporarily redirect stdout or stderr

    e.g.:

    with stdchannel_redirected(sys.stderr, os.devnull):
        if compiler.has_function('clock_gettime', libraries=['rt']):
            libraries.append('rt')

    Code adapted from https://stackoverflow.com/a/17752455/1382869
    """

    try:
        oldstdchannel = os.dup(stdchannel.fileno())
        dest_file = open(dest_filename, 'w')
        os.dup2(dest_file.fileno(), stdchannel.fileno())

        yield
    finally:
        if oldstdchannel is not None:
            os.dup2(oldstdchannel, stdchannel.fileno())
        if dest_file is not None:
            dest_file.close()

def check_for_openmp():
    """Returns True if local setup supports OpenMP, False otherwise

    Code adapted from astropy_helpers, originally written by Tom 
    Robitaille and Curtis McCully.
    """

    # See https://bugs.python.org/issue25150
    if sys.version_info[:3] == (3, 5, 0):
        return False

    # Create a temporary directory
    ccompiler = new_compiler()
    customize_compiler(ccompiler)

    tmp_dir = tempfile.mkdtemp()
    start_dir = os.path.abspath('.')

    if os.name == 'nt':
        # TODO: make this work with mingw
        # AFAICS there's no easy way to get the compiler distutils
        # will be using until compilation actually happens
        compile_flag = '-openmp'
        link_flag = ''
    else:
        compile_flag = '-fopenmp'
        link_flag = '-fopenmp'

    try:
        os.chdir(tmp_dir)

        with open('test_openmp.c', 'w') as f:
            f.write(CCODE)

        os.mkdir('objects')

        # Compile, link, and run test program
        with stdchannel_redirected(sys.stderr, os.devnull):
            ccompiler.compile(['test_openmp.c'], output_dir='objects',
                              extra_postargs=[compile_flag])
            ccompiler.link_executable(
                glob.glob(os.path.join('objects', '*')), 'test_openmp',
                extra_postargs=[link_flag])
            output = subprocess.check_output('./test_openmp').decode(
                sys.stdout.encoding or 'utf-8').splitlines()

        if 'nthreads=' in output[0]:
            nthreads = int(output[0].strip().split('=')[1])
            if len(output) == nthreads:
                using_openmp = True
            else:
                log.warn("Unexpected number of lines from output of test "
                         "OpenMP program (output was {0})".format(output))
                using_openmp = False
        else:
            log.warn("Unexpected output from test OpenMP "
                     "program (output was {0})".format(output))
            using_openmp = False

    except (CompileError, LinkError):
        using_openmp = False
    finally:
        os.chdir(start_dir)

    if using_openmp:
        log.warn("Using OpenMP to compile parallel extensions")
    else:
        log.warn("Unable to compile OpenMP test program so Cython\n"
                 "extensions will be compiled without parallel support")

    return using_openmp


def check_for_pyembree():
    try:
        fn = resource_filename("pyembree", "rtcore.pxd")
    except ImportError:
        return None
    return os.path.dirname(fn)

def in_conda_env():
    return any(s in sys.version for s in ("Anaconda", "Continuum"))

def read_embree_location():
    '''

    Attempts to locate the embree installation. First, we check for an
    EMBREE_DIR environment variable. If one is not defined, we look for
    an embree.cfg file in the root yt source directory. Finally, if that
    is not present, we default to /usr/local. If embree is installed in a
    non-standard location and none of the above are set, the compile will
    not succeed. This only gets called if check_for_pyembree() returns
    something other than None.

    '''

    rd = os.environ.get('EMBREE_DIR')
    if rd is None:
        try:
            rd = open("embree.cfg").read().strip()
        except IOError:
            rd = '/usr/local'

    fail_msg = ("I attempted to find Embree headers in %s. \n"
               "If this is not correct, please set your correct embree location \n"
               "using EMBREE_DIR environment variable or your embree.cfg file. \n"
               "Please see http://yt-project.org/docs/dev/visualizing/unstructured_mesh_rendering.html "
                "for more information. \n" % rd)

    # Create a temporary directory
    tmpdir = tempfile.mkdtemp()
    curdir = os.getcwd()

    try:
        os.chdir(tmpdir)

        # Get compiler invocation
        compiler = os.getenv('CXX', 'c++')
        compiler = compiler.split(' ')

        # Attempt to compile a test script.
        filename = r'test.cpp'
        file = open(filename, 'wt', 1)
        file.write(
            '#include "embree2/rtcore.h"\n'
            'int main() {\n'
            'return 0;\n'
            '}'
        )
        file.flush()
        p = Popen(compiler + ['-I%s/include/' % rd, filename], 
                  stdin=PIPE, stdout=PIPE, stderr=PIPE)
        output, err = p.communicate()
        exit_code = p.returncode

        if exit_code != 0:
            log.warn("Pyembree is installed, but I could not compile Embree "
                     "test code.")
            log.warn("The error message was: ")
            log.warn(err)
            log.warn(fail_msg)

        # Clean up
        file.close()

    except OSError:
        log.warn("read_embree_location() could not find your C compiler. "
                 "Attempted to use '%s'. " % compiler)
        return False

    finally:
        os.chdir(curdir)
        shutil.rmtree(tmpdir)

    return rd