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