Coverage for mango/imports.py : 58%
Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1from importlib import util as imp_u, import_module as imp
2from os import listdir
3from mango.constants import c
6def read(save_type):
7 """
8 Import modules based on options for reading data files.
10 Parameters
11 ----------
12 save_type: str
13 file type extension desired
15 Returns
16 -------
17 load: func
18 loading files function
19 save_type: str
20 save type based on available
22 """
23 if save_type == "hdf5":
24 try:
25 from mango.hdf5io import load
26 new_save_type = 'hdf5'
27 except ImportError:
28 from pickle import load
29 c.Error("M Using pickle for data reading\n install pytables for hdf5 r/w")
30 new_save_type = "pkl"
31 elif save_type == "pkl":
32 from pickle import load
33 new_save_type = "pkl"
34 elif save_type == 'xyz':
35 load = None
36 new_save_type = 'xyz'
37 else:
38 new_save_type = "txt"
39 load = None
40 if save_type != new_save_type: 40 ↛ 41line 40 didn't jump to line 41, because the condition on line 40 was never true
41 c.Error(f'W Unable to load {save_type} reader')
42 return load, save_type
45def write(save_type):
46 """
47 Import modules based on options for writing data files.
49 Parameters
50 ----------
51 save_type: str
52 file type extension desired
54 Returns
55 -------
56 dump: func
57 save files function
58 save_type: str
59 save type based on available
61 """
62 if save_type == "hdf5":
63 try:
64 from mango.hdf5io import save as dump
65 except ImportError:
66 from pickle import dump
67 c.Error("M Using pickle for data reading\n install pytables for hdf5 r/w")
68 save_type = "pkl"
69 elif save_type == "pkl":
70 from pickle import dump
71 else:
72 save_type = "txt"
73 dump = None
74 return dump, save_type
77def inquirer():
78 """
79 Import inquirer if required.
81 Returns
82 -------
83 bool:
84 Is inquirere available?
85 tuple:
86 inquirer functions
88 """
89 if c.opsys != "nt" and c.tinfo['otty']:
90 try:
91 from mango.pp.util import captured_output
92 with captured_output() as (out, err):
93 from inquirer import Checkbox
94 from inquirer.render.console import ConsoleRender
95 return True, (Checkbox, ConsoleRender)
96 except ImportError:
97 return False, (None, None)
100class matplotlib():
101 """Importing matplotlib when necessary, allows for changing the backend before import."""
103 def __init__(self, backend=None):
104 from matplotlib import use, rcParams
105 if c.havedisplay is False: 105 ↛ 106line 105 didn't jump to line 106, because the condition on line 105 was never true
106 use('Agg')
107 rcParams['agg.path.chunksize'] = 10000
108 elif backend is not None: 108 ↛ 109line 108 didn't jump to line 109, because the condition on line 108 was never true
109 use(backend)
110 from matplotlib.pyplot import figure, savefig
112 self.figure = figure
113 self.savefig = savefig
114 self.rcParams = rcParams
116 def suscep(self):
117 from matplotlib.lines import Line2D
118 from matplotlib.pyplot import legend, yscale
119 return (self.figure, self.savefig, self.errorbarpicker(),
120 Line2D, legend, yscale)
122 def col(self):
123 return self.figure, self.savefig, self.htuple()
125 @staticmethod
126 def showg():
127 from matplotlib.pyplot import show, get_fignums
128 return show, get_fignums
130 @staticmethod
131 def close():
132 from matplotlib.pyplot import close
133 return close
135 @staticmethod
136 def errorbarpicker():
137 import matplotlib.legend_handler as mlh
138 from matplotlib.container import ErrorbarContainer
140 class newHandler(mlh.HandlerErrorbar):
141 def create_artists(self, *args, **kwargs):
142 a_list = mlh.HandlerErrorbar.create_artists(self, *args, **kwargs)
143 a_list = a_list[-2:] + a_list[:-2]
144 return a_list
146 return {ErrorbarContainer: newHandler()}
148 def prettyplot(self):
149 from matplotlib import ticker, __version__ as mplversion
150 from matplotlib.pyplot import setp, gcf
151 from matplotlib.backends.backend_pdf import PdfPages
152 from subprocess import getstatusoutput
153 if getstatusoutput('which latex')[0] == 0:
154 self.rcParams['text.usetex'] = True
155 self.rcParams['font.family'] = 'MathJax_Main'
156 self.rcParams['font.size'] = 14
157 if int(mplversion[0]) < 3:
158 self.rcParams['text.latex.unicode'] = True
160 return self.figure, PdfPages, setp, gcf, ticker, self.htuple()
162 def htuple(self):
163 """single legend entry for multiple connected graphs with one legend. see mango.tools.plotting"""
164 import matplotlib.legend_handler as lh
165 import numpy as np
166 from matplotlib.collections import PolyCollection
168 class mytuple(lh.HandlerTuple):
169 def create_artists(self, legend, orig_handle,
170 xdescent, ydescent, width, height, fontsize,
171 trans):
173 handler_map = legend.get_legend_handler_map()
175 if self._ndivide is None:
176 ndivide = len(orig_handle)
177 else:
178 ndivide = self._ndivide
180 if self._pad is None:
181 pad = legend.borderpad * fontsize
182 else:
183 pad = self._pad * fontsize
185 if ndivide > 1:
186 width = (width - pad * (ndivide - 1)) / ndivide
188 xds_cycle = lh.cycle(xdescent - (width + pad) * np.arange(ndivide))
190 a_list = []
191 c_list = [0, 1] if isinstance(orig_handle[0], PolyCollection) else [0]
192 for no, handle1 in enumerate(orig_handle):
193 handler = legend.get_legend_handler(handler_map, handle1)
194 if no in c_list:
195 _a_list = handler.create_artists(
196 legend, handle1,
197 next(xds_cycle), ydescent, width, height, fontsize, trans)
198 a_list.extend(_a_list)
200 return a_list
201 return {tuple: mytuple()}
204def fftw():
205 """
206 Import the fastest version of fftw installed.
208 Returns
209 -------
210 output: fft, ffts, ifft
211 Fastest fft implementations available
213 """
214 from scipy import fftpack as ffts
215 try:
216 from pyfftw.interfaces import scipy_fftpack
217 from pyfftw import byte_align
218 fft = scipy_fftpack.fft
219 ifft = scipy_fftpack.ifft
221 except ImportError:
222 def fft(arr, n=None, axis=-1, threads=None):
223 return ffts.fft(arr, n=n, axis=axis)
225 def ifft(arr, n=None, axis=-1, threads=None):
226 return ffts.ifft(arr, n=n, axis=axis)
228 def byte_align(array, n=None, dtype=None):
229 return array
231 c.Error(">M Using scipy fft for Susceptibility\n install fftw and pyfftw for better performance")
232 return fft, ffts, ifft, byte_align
235def getpotentials(name):
236 """
237 Load pluggable potentials.
239 Parameters
240 ----------
241 name: string
242 Filename of potentials module to import
244 Returns
245 -------
246 output: module
247 Loaded module
249 """
250 try:
251 module = imp("mango.potentials.{}".format(name))
252 except ImportError:
253 module = _loadfromspec(name)
254 c.Error(f"M Loaded {module.__name__}")
255 try:
256 for i in c.posfuncs:
257 if i not in module.__all__:
258 c.Error("W function {} not replaced.".format(i))
259 except AttributeError:
260 c.Error("F module has no attribute __all__")
261 return module
264def _loadfromspec(name):
265 full_dirname = name.rsplit("/", 1)
266 dirname = '.' if len(full_dirname[0]) == 0 else full_dirname[0]
268 try:
269 mod_files = [file for file in listdir(dirname) if file.startswith(full_dirname[1])]
270 except FileNotFoundError:
271 c.Error("F Can't find module file '{}'".format(name))
273 requested = full_dirname[1] if full_dirname[1] in mod_files else mod_files[0]
275 if len(mod_files) > 1: 275 ↛ 276line 275 didn't jump to line 276, because the condition on line 275 was never true
276 c.Error("{}{}".format("W Multiple files start with '{}'\n".format(full_dirname[1]),
277 "Assuming module is '{}'".format(requested)))
279 mod_file = f"{dirname}/{requested}"
281 try:
282 spec = imp_u.spec_from_file_location(mod_file.rsplit("/")[-1].split(".")[0], mod_file)
283 module = imp_u.module_from_spec(spec)
284 spec.loader.exec_module(module)
285 except (AttributeError, ImportError):
286 c.Error("F File '{}' is not a module".format(mod_files[0]))
288 return module