From 9ad9b00207c75236f4e24984717a323160d7160d Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Tue, 19 Mar 2024 15:55:49 -0500
Subject: [PATCH] Compiler cache: honor XDG_CACHE_HOME on MacOS

---
 pyopencl/cache.py  | 22 +++++++++++-----------
 test/test_array.py | 44 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 11 deletions(-)

diff --git a/pyopencl/cache.py b/pyopencl/cache.py
index 96906137f..a44c6a38c 100644
--- a/pyopencl/cache.py
+++ b/pyopencl/cache.py
@@ -348,20 +348,20 @@ def _create_built_program_from_source_cached(ctx, src, options_bytes,
         except ImportError:
             import appdirs
 
-        cache_dir = join(appdirs.user_cache_dir("pyopencl", "pyopencl"),
+        # Determine the cache directory in the same way as pytools.PersistentDict,
+        # which PyOpenCL uses for invoker caches.
+        if sys.platform == "darwin" and os.getenv("XDG_CACHE_HOME") is not None:
+            # appdirs and platformdirs do not handle XDG_CACHE_HOME on macOS
+            # https://github.com/platformdirs/platformdirs/issues/269
+            cache_dir = join(os.getenv("XDG_CACHE_HOME"), "pyopencl")
+        else:
+            cache_dir = appdirs.user_cache_dir("pyopencl", "pyopencl")
+
+        cache_dir = join(cache_dir,
                 "pyopencl-compiler-cache-v2-py{}".format(
                     ".".join(str(i) for i in sys.version_info)))
 
-    # {{{ ensure cache directory exists
-
-    try:
-        os.makedirs(cache_dir)
-    except OSError as e:
-        from errno import EEXIST
-        if e.errno != EEXIST:
-            raise
-
-    # }}}
+    os.makedirs(cache_dir, exist_ok=True)
 
     if devices is None:
         devices = ctx.devices
diff --git a/test/test_array.py b/test/test_array.py
index 95fe4e74a..2cc62e8f8 100644
--- a/test/test_array.py
+++ b/test/test_array.py
@@ -2350,6 +2350,50 @@ def test_logical_not(ctx_factory):
         np.logical_not(np.ones(10)))
 
 
+# {{{ test XDG_CACHE_HOME handling
+
+@pytest.mark.skipif(sys.platform == "win32",
+                    reason="XDG_CACHE_HOME is not used on Windows")
+def test_xdg_cache_home(ctx_factory):
+    import os
+    import shutil
+    from os.path import join
+
+    context = ctx_factory()
+    queue = cl.CommandQueue(context)
+
+    a = np.array([1, 2, 3, 4, 5]).astype(np.float32)
+    a_gpu = cl_array.to_device(queue, a)
+
+    xdg_dir = "tmpdir_pyopencl_xdg_test"
+
+    # PyOpenCL uses pytools.PersistentDict for invoker caches,
+    # which is why xdg_dir will always exist. Therefore, check
+    # whether xdg_pyopencl_dir exists.
+    xdg_pyopencl_dir = join(xdg_dir, "pyopencl")
+    assert not os.path.exists(xdg_dir)
+
+    old_xdg_cache_home = None
+
+    try:
+        old_xdg_cache_home = os.getenv("XDG_CACHE_HOME")
+        os.environ["XDG_CACHE_HOME"] = xdg_dir
+
+        result = pow(a_gpu, a_gpu).get()
+        assert (np.abs(a ** a - result) < 3e-3).all()
+
+        assert os.path.exists(xdg_pyopencl_dir)
+    finally:
+        if old_xdg_cache_home is not None:
+            os.environ["XDG_CACHE_HOME"] = old_xdg_cache_home
+        else:
+            del os.environ["XDG_CACHE_HOME"]
+
+        shutil.rmtree(xdg_dir)
+
+# }}}
+
+
 if __name__ == "__main__":
     if len(sys.argv) > 1:
         exec(sys.argv[1])