Use add_dll_directory as a context manager

Author: maresb

pytensor/link/c/cmodule.py

@@ -4,6 +4,7 @@
 """
 
 import atexit
+from contextlib import nullcontext
 import importlib
 import logging
 import os
@@ -23,7 +24,7 @@
 from collections.abc import Callable
 from functools import cache
 from io import BytesIO, StringIO
-from typing import TYPE_CHECKING, Protocol, cast
+from typing import TYPE_CHECKING, Generator, Protocol, cast
 
 import numpy as np
 from setuptools._distutils.sysconfig import (
@@ -272,14 +273,13 @@ def _get_ext_suffix():
     return dist_suffix
 
 
-@cache  # See explanation in docstring.
-def add_gcc_dll_directory() -> None:
+def add_gcc_dll_directory() -> Generator[None, None, None]:
     """On Windows, detect and add the location of gcc to the DLL search directory.
 
     On non-Windows platforms this is a noop.
 
-    The @cache decorator ensures that this function only executes once to avoid
-    redundant entries. See <https://github.com/pymc-devs/pytensor/pull/678>.
+    Returns a context manager to be used with `with`. The entry is removed when the
+    context manager is closed. See <https://github.com/pymc-devs/pytensor/pull/678>.
     """
     if (sys.platform == "win32") & (hasattr(os, "add_dll_directory")):
         gcc_path = shutil.which("gcc")
@@ -288,7 +288,8 @@ def add_gcc_dll_directory() -> None:
             # the ignore[attr-defined] on non-Windows platforms.
             # For Windows we need ignore[unused-ignore] since the ignore
             # is unnecessary with that platform.
-            os.add_dll_directory(os.path.dirname(gcc_path))  # type: ignore[attr-defined,unused-ignore]
+            return os.add_dll_directory(os.path.dirname(gcc_path))  # type: ignore[attr-defined,unused-ignore]
+    return nullcontext()
 
 
 def dlimport(fullpath, suffix=None):
@@ -340,20 +341,20 @@ def dlimport(fullpath, suffix=None):
     _logger.debug(f"module_name {module_name}")
 
     sys.path[0:0] = [workdir]  # insert workdir at beginning (temporarily)
-    add_gcc_dll_directory()
-    global import_time
-    try:
-        importlib.invalidate_caches()
-        t0 = time.perf_counter()
-        with warnings.catch_warnings():
-            warnings.filterwarnings("ignore", message="numpy.ndarray size changed")
-            rval = __import__(module_name, {}, {}, [module_name])
-        t1 = time.perf_counter()
-        import_time += t1 - t0
-        if not rval:
-            raise Exception("__import__ failed", fullpath)
-    finally:
-        del sys.path[0]
+    with add_gcc_dll_directory():
+        global import_time
+        try:
+            importlib.invalidate_caches()
+            t0 = time.perf_counter()
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", message="numpy.ndarray size changed")
+                rval = __import__(module_name, {}, {}, [module_name])
+            t1 = time.perf_counter()
+            import_time += t1 - t0
+            if not rval:
+                raise Exception("__import__ failed", fullpath)
+        finally:
+            del sys.path[0]
 
     assert fullpath.startswith(rval.__file__)
     return rval