From 2a6c08f921bc8e58bd3fb9d8cda30a68f76b758a Mon Sep 17 00:00:00 2001
From: Ryan Stelzleni <rstelzleni@nvidia.com>
Date: Fri, 30 Apr 2021 16:25:55 -0700
Subject: [PATCH] Support for Python 3.8+ on Windows

Previously this had been disabled because of a change to how the python
interpreter searches for DLLs. This change should restore the old
behavior, and give software builders the ability to opt into the new
python behavior where possible. See

https://docs.python.org/3/whatsnew/3.8.html#bpo-36085-whatsnew
---
 BUILDING.md                |  7 +++++++
 build_scripts/build_usd.py |  9 --------
 pxr/base/tf/__init__.py    | 43 +++++++++++++++++++++++++++++++++++++-
 3 files changed, 49 insertions(+), 10 deletions(-)

diff --git a/BUILDING.md b/BUILDING.md
index 3f50ba4d2d..0d9781c919 100644
--- a/BUILDING.md
+++ b/BUILDING.md
@@ -603,3 +603,10 @@ platforms to avoid issues with Boost config files (introduced in Boost version
 to use Boost specified config files for their USD build, specify 
 -DBoost_NO_BOOST_CMAKE=OFF when running cmake.
 
+2. Windows and Python 3.8+
+Python 3.8 and later on Windows will no longer search PATH for DLL dependencies.
+Instead, clients can call `os.add_dll_directory(p)` to set paths to search.
+By default on that platform USD will iterate over PATH and add all paths using
+`os.add_dll_directory()` when importing Python modules. Users may override
+this by setting the environment variable `PXR_USD_WINDOWS_DLL_PATH` to a PATH-like
+string. If this is set, USD will use these paths instead.
diff --git a/build_scripts/build_usd.py b/build_scripts/build_usd.py
index e94598a920..d294290c20 100644
--- a/build_scripts/build_usd.py
+++ b/build_scripts/build_usd.py
@@ -1964,15 +1964,6 @@ def ForceBuildDependency(self, dep):
                "PATH")
     sys.exit(1)
 
-# Error out on Windows with Python 3.8+. USD currently does not support
-# these versions due to:
-# https://docs.python.org/3.8/whatsnew/3.8.html#bpo-36085-whatsnew
-isPython38 = (sys.version_info.major >= 3 and
-              sys.version_info.minor >= 8)
-if Windows() and isPython38:
-    PrintError("Python 3.8+ is not supported on Windows")
-    sys.exit(1)
-
 if find_executable("cmake"):
     # Check cmake requirements
     if Windows():
diff --git a/pxr/base/tf/__init__.py b/pxr/base/tf/__init__.py
index 0a324abe14..53e437642d 100644
--- a/pxr/base/tf/__init__.py
+++ b/pxr/base/tf/__init__.py
@@ -25,6 +25,44 @@
 Tf -- Tools Foundation
 """
 
+# Type to help handle DLL import paths on Windows with python interpreters v3.8
+# and newer. These interpreters don't search for DLLs in the path anymore, you
+# have to provide a path explicitly. This re-enables path searching for USD 
+# dependency libraries
+import platform, sys
+if sys.version_info >= (3, 8) and platform.system() == "Windows":
+    import contextlib
+
+    @contextlib.contextmanager
+    def WindowsImportWrapper():
+        import os
+        dirs = []
+        import_paths = os.getenv('PXR_USD_WINDOWS_DLL_PATH')
+        if import_paths is None:
+            import_paths = os.getenv('PATH', '')
+        for path in import_paths.split(os.pathsep):
+            # Calling add_dll_directory raises an exception if paths don't
+            # exist, or if you pass in dot
+            if os.path.exists(path) and path != '.':
+                dirs.append(os.add_dll_directory(path))
+        # This block guarantees we clear the dll directories if an exception
+        # is raised in the with block.
+        try:
+            yield
+        finally:
+            for dll_dir in dirs:
+                dll_dir.close()
+        del os
+    del contextlib
+else:
+    class WindowsImportWrapper(object):
+        def __enter__(self):
+            pass
+        def __exit__(self, exc_type, ex_val, exc_tb):
+            pass
+del platform, sys
+
+
 def PreparePythonModule(moduleName=None):
     """Prepare an extension module at import time.  This will import the
     Python module associated with the caller's module (e.g. '_tf' for 'pxr.Tf')
@@ -46,7 +84,10 @@ def PreparePythonModule(moduleName=None):
             moduleName = f_locals["__name__"].split(".")[-1]
             moduleName = "_" + moduleName[0].lower() + moduleName[1:]
 
-        module = importlib.import_module("." + moduleName, f_locals["__name__"])
+        with WindowsImportWrapper():
+            module = importlib.import_module(
+                    "." + moduleName, f_locals["__name__"])
+
         PrepareModule(module, f_locals)
         try:
             del f_locals[moduleName]