1"""Global configuration variables for PyFFTW.
2
3The approach taken here was adapated from Numba's config.py.
4"""
5from __future__ import print_function, division, absolute_import
6
7import os
8import multiprocessing
9import warnings
10
11from .pyfftw import _threading_type
12
13
14class _EnvReloader(object):
15
16    def __init__(self):
17        self.reset()
18
19    def reset(self):
20        self.old_environ = {}
21        self.update(force=True)
22
23    def update(self, force=False):
24        new_environ = {}
25
26        # read local env var OMP_NUM_THREADS and any starting with PYFFTW_
27        for name, value in os.environ.items():
28            if name.startswith('PYFFTW_') or name == 'OMP_NUM_THREADS':
29                new_environ[name] = value
30        # We update the config variables if at least one PYFFTW environment
31        # variable was modified.  This lets the user modify values
32        # directly in the config module without having them when
33        # reload_config() is called by the compiler.
34        if force or self.old_environ != new_environ:
35            self.process_environ(new_environ)
36            # Store a copy
37            self.old_environ = dict(new_environ)
38
39    def process_environ(self, environ):
40        def _readenv(name, ctor, default):
41            value = environ.get(name)
42            if value is None:
43                return default() if callable(default) else default
44            try:
45                return ctor(value)
46            except Exception:
47                warnings.warn("environ %s defined but failed to parse '%s'" %
48                              (name, value), RuntimeWarning)
49                return default
50
51        def optional_str(x):
52            return str(x) if x is not None else None
53
54        if _threading_type is None:
55            NUM_THREADS = 1
56        else:
57            if (_threading_type == "OMP" and
58                    "PYFFTW_NUM_THREADS" not in environ):
59                # fallback to OMP_NUM_THREADS if PYFFTW_NUM_THREADS undefined
60                NUM_THREADS = _readenv("OMP_NUM_THREADS", int, 1)
61            else:
62                NUM_THREADS = _readenv("PYFFTW_NUM_THREADS", int, 1)
63            # if user requested <= 0 threads, use the maximum available
64            if NUM_THREADS <= 0:
65                NUM_THREADS = multiprocessing.cpu_count()
66
67        PLANNER_EFFORT = _readenv(
68            "PYFFTW_PLANNER_EFFORT", str, "FFTW_ESTIMATE")
69
70        # Inject the configuration values into the module globals
71        for name, value in locals().copy().items():
72            if name.isupper():
73                globals()[name] = value
74
75_env_reloader = _EnvReloader()
76
77
78def _reload_config():
79    """
80    Reload the configuration from environment variables, if necessary.
81    """
82    _env_reloader.update()
83