Add QUATERNIONIC_DISABLE_CACHE environment variable.

It may not always be desirable to cache the Numba-generated code, for
example in a HPC environment where having lots of workers each compiling
the code may be preferable to them all trying to load it from the cache.

This adds a QUATERNIONIC_DISABLE_CACHE environment variable. If set to
1, the `cache` parameter of the `jit` and `guvectorize` decorators is
set to False. Default behaviour remains the same as before, i.e., to
use the cache.

README.md

@@ -273,6 +273,15 @@ Finally, there are also capabilities related to interpolation, for example as fu
   * squad (spherical quadratic interpolation)
 
 
+## Caching
+
+By default, the compiled code generated by numba is cached so that the compilation only needs to
+take place on the first import. If you want to disable this caching, for example in a
+high-performance computing environment where it may be preferable to compile the code than try to
+load a cache from disk, set the environment variable `QUATERNIONIC_DISABLE_CACHE` to `1` before
+importing this package.
+
+
 # Related packages
 
 Other python packages with some quaternion features include

quaternionic/utilities.py

@@ -2,6 +2,7 @@
 # See LICENSE file for details:
 # <https://github.com/moble/quaternionic/blob/master/LICENSE>
 
+import os
 import sys
 import functools
 import numpy as np
@@ -213,5 +214,6 @@ def __getitem__(self, *args, **kwargs):
 else:
     import numba
     from numba import float64, boolean
-    jit = functools.partial(numba.njit, cache=True)
-    guvectorize = functools.partial(numba.guvectorize, nopython=True, cache=True)
+    cache = os.environ.get("QUATERNIONIC_DISABLE_CACHE", "0") != "1"
+    jit = functools.partial(numba.njit, cache=cache)
+    guvectorize = functools.partial(numba.guvectorize, nopython=True, cache=cache)

tests/test_utilities.py

@@ -1,3 +1,4 @@
+import subprocess
 import sys
 import numpy as np
 import quaternionic
@@ -118,3 +119,37 @@ def test_pyguvectorize():
                     atol=0.0,
                     rtol=_quaternion_resolution
                 )
+
+
+@pytest.mark.skipif(sys.implementation.name.lower() == "pypy", reason="No numba on pypy")
+def test_cache_disable(tmp_path):
+    # First check caching works by default.
+    cache_dir = tmp_path / "enabled"
+    subprocess.run(
+        [sys.executable, "-c", "import quaternionic"],
+        env={
+            "NUMBA_CACHE_DIR": str(cache_dir),
+        },
+    )
+
+    # Numba uses a subdirectory named `quaternionic_<hashstr>`, with the hashstr computed
+    # from the full path to the source. We should have 1 directory in our cache, and that
+    # should contain many files.
+    contents = list(cache_dir.iterdir())
+    assert len(contents) == 1
+    subdir = contents[0]
+    assert subdir.name.startswith("quaternionic_")
+    assert len(list(subdir.iterdir())) > 10
+
+    # Change the cache location and check with the environment variable disabling caching.
+    # This seems to not create the cache directory, but include an alternative check that
+    # it is empty if it exists in case the directory does get created sometimes.
+    cache_dir = tmp_path / "disabled"
+    subprocess.run(
+        [sys.executable, "-c", "import quaternionic"],
+        env={
+            "QUATERNIONIC_DISABLE_CACHE": "1",
+            "NUMBA_CACHE_DIR": str(cache_dir),
+        },
+    )
+    assert (not cache_dir.exists()) or len(list(cache_dir.iterdir())) == 0