From 1696fba31e3626097690882ea96d7235738b67d8 Mon Sep 17 00:00:00 2001
From: Louis Montaut <louismontaut@gmail.com>
Date: Mon, 24 Jun 2024 15:21:49 +0200
Subject: [PATCH] python: fix windows build

---
 python/CMakeLists.txt              |  3 ++
 python/coal/__init__.py            | 29 ++++++++++++-
 python/coal/windows_dll_manager.py | 65 ++++++++++++++++++++++++++++++
 3 files changed, 95 insertions(+), 2 deletions(-)
 create mode 100644 python/coal/windows_dll_manager.py

diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt
index cab976533..ca6e00e13 100644
--- a/python/CMakeLists.txt
+++ b/python/CMakeLists.txt
@@ -142,6 +142,8 @@ SET_TARGET_PROPERTIES(${PYTHON_LIB_NAME} PROPERTIES
   SUFFIX "${PYTHON_EXT_SUFFIX}"
   OUTPUT_NAME "${PYTHON_LIB_NAME}"
   LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/python/${PROJECT_NAME}"
+  # On Windows, shared library are treat as binary
+  RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/python/${PROJECT_NAME}"
   )
 
 IF(IS_ABSOLUTE ${PYTHON_SITELIB})
@@ -169,6 +171,7 @@ ENDIF(GENERATE_PYTHON_STUBS)
 SET(PYTHON_FILES
   __init__.py
   viewer.py
+  windows_dll_manager.py
   )
 
 FOREACH(python ${PYTHON_FILES})
diff --git a/python/coal/__init__.py b/python/coal/__init__.py
index e26c55ab9..966642738 100644
--- a/python/coal/__init__.py
+++ b/python/coal/__init__.py
@@ -32,5 +32,30 @@
 #  POSSIBILITY OF SUCH DAMAGE.
 
 # ruff: noqa: F401, F403
-from .coal_pywrap import *
-from .coal_pywrap import __raw_version__, __version__
+
+# On Windows, if coal.dll is not in the same directory than
+# the .pyd, it will not be loaded.
+# We first try to load coal, then, if it fail and we are on Windows:
+#  1. We add all paths inside COAL_WINDOWS_DLL_PATH to DllDirectory
+#  2. If COAL_WINDOWS_DLL_PATH we add the relative path from the
+#     package directory to the bin directory to DllDirectory
+# This solution is inspired from:
+#  - https://github.com/PixarAnimationStudios/OpenUSD/pull/1511/files
+#  - https://stackoverflow.com/questions/65334494/python-c-extension-packaging-dll-along-with-pyd
+# More resources on https://github.com/diffpy/pyobjcryst/issues/33
+try:
+    from .coal_pywrap import *  # noqa
+    from .coal_pywrap import __raw_version__, __version__
+except ImportError:
+    import platform
+
+    if platform.system() == "Windows":
+        from .windows_dll_manager import build_directory_manager, get_dll_paths
+
+        with build_directory_manager() as dll_dir_manager:
+            for p in get_dll_paths():
+                dll_dir_manager.add_dll_directory(p)
+            from .coal_pywrap import *  # noqa
+            from .coal_pywrap import __raw_version__, __version__  # noqa
+    else:
+        raise
diff --git a/python/coal/windows_dll_manager.py b/python/coal/windows_dll_manager.py
new file mode 100644
index 000000000..92516fee1
--- /dev/null
+++ b/python/coal/windows_dll_manager.py
@@ -0,0 +1,65 @@
+import contextlib
+import os
+import sys
+
+
+def get_dll_paths():
+    coal_paths = os.getenv("COAL_WINDOWS_DLL_PATH")
+    if coal_paths is None:
+        # From https://peps.python.org/pep-0250/#implementation
+        # lib/python-version/site-packages/package
+        RELATIVE_DLL_PATH1 = "..\\..\\..\\..\\bin"
+        # lib/site-packages/package
+        RELATIVE_DLL_PATH2 = "..\\..\\..\\bin"
+        # For unit test
+        RELATIVE_DLL_PATH3 = "..\\..\\bin"
+        return [
+            os.path.join(os.path.dirname(__file__), RELATIVE_DLL_PATH1),
+            os.path.join(os.path.dirname(__file__), RELATIVE_DLL_PATH2),
+            os.path.join(os.path.dirname(__file__), RELATIVE_DLL_PATH3),
+        ]
+    else:
+        return coal_paths.split(os.pathsep)
+
+
+class PathManager(contextlib.AbstractContextManager):
+    """Restore PATH state after importing Python module"""
+
+    def add_dll_directory(self, dll_dir: str):
+        os.environ["PATH"] += os.pathsep + dll_dir
+
+    def __enter__(self):
+        self.old_path = os.environ["PATH"]
+        return self
+
+    def __exit__(self, *exc_details):
+        os.environ["PATH"] = self.old_path
+
+
+class DllDirectoryManager(contextlib.AbstractContextManager):
+    """Restore DllDirectory state after importing Python module"""
+
+    def add_dll_directory(self, dll_dir: str):
+        # add_dll_directory can fail on relative path and non
+        # existing path.
+        # Since we don't know all the fail criterion we just ignore
+        # thrown exception
+        try:
+            self.dll_dirs.append(os.add_dll_directory(dll_dir))
+        except OSError:
+            pass
+
+    def __enter__(self):
+        self.dll_dirs = []
+        return self
+
+    def __exit__(self, *exc_details):
+        for d in self.dll_dirs:
+            d.close()
+
+
+def build_directory_manager():
+    if sys.version_info >= (3, 8):
+        return DllDirectoryManager()
+    else:
+        return PathManager()