implement mark_jit_unsafe(cache=False)

CalebBell · Feb 28, 2021 · 57e3212 · 57e3212
1 parent 33b6f03
commit 57e3212
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 8 deletions.
diff --git a/chemicals/numba.py b/chemicals/numba.py
@@ -17,7 +17,7 @@
 TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR
 THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 """
-from chemicals.utils import PY37, numba_blacklisted
+from chemicals.utils import PY37, numba_blacklisted, numba_cache_blacklisted
 
 busy = False
 __all__ = []
@@ -33,11 +33,14 @@ def transform():
     normal_fluids = fluids
     orig_file = __file__
     def transform_complete_chemicals(replaced, __funcs, __all__, normal, vec):
-        cache_blacklist = set( ['Rachford_Rice_solution', 'Rachford_Rice_solution_LN2', 'Rachford_Rice_solution_polynomial',
-                  'Rachford_Rice_solution_numpy', 'Li_Johns_Ahmadi_solution', 'flash_inner_loop',
-                  'Rachford_Rice_solutionN', 'Rachford_Rice_solution2', 'flash_wilson',
-                  'Lastovka_Shaw_T_for_Hm', 'Lastovka_Shaw_T_for_Sm', 'iapws97_T', 'iapws95_T'])
-
+        numba_cache_blacklisted.extend([
+            'Rachford_Rice_solution', 'Rachford_Rice_solution_LN2',
+            'Rachford_Rice_solution_polynomial','Rachford_Rice_solution_numpy', 
+            'Li_Johns_Ahmadi_solution', 'flash_inner_loop',
+            'Rachford_Rice_solutionN', 'Rachford_Rice_solution2', 'flash_wilson',
+            'Lastovka_Shaw_T_for_Hm', 'Lastovka_Shaw_T_for_Sm', 'iapws97_T', 
+            'iapws95_T'])
+        cache_blacklist = set(numba_cache_blacklisted)
         __funcs.update(normal_fluids.numba.numbafied_fluids_functions.copy())
         new_mods = normal_fluids.numba.transform_module(normal, __funcs, replaced, vec=vec,
                                                         blacklist=set(numba_blacklisted),

diff --git a/chemicals/utils.py b/chemicals/utils.py
@@ -96,8 +96,11 @@ def os_path_join(*args):
     pass
 
 numba_blacklisted = []
-def mark_jit_unsafe(f):
-    numba_blacklisted.append(f.__name__)
+numba_cache_blacklisted = []
+
+def mark_jit_unsafe(f=None, cache=False):
+    if not f: return lambda f: mark_jit_unsafe(f, cache)
+    (numba_cache_blacklisted if cache else numba_blacklisted).append(f.__name__)
     return f
 
 @mark_jit_unsafe