More review feedback addressed

pytorch · Dec 13, 2021 · 2d8c032 · 2d8c032
1 parent d6129e9
commit 2d8c032
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 20 deletions.
diff --git a/torchvision/_internally_replaced_utils.py b/torchvision/_internally_replaced_utils.py
@@ -22,7 +22,7 @@ def _is_remote_location_available() -> bool:
     from torch.utils.model_zoo import load_url as load_state_dict_from_url  # noqa: 401
 
 
-def _get_extension_path(lib_name) -> str:
+def _get_extension_path(lib_name: str) -> str:
 
     lib_dir = os.path.dirname(__file__)
     if os.name == "nt":
@@ -56,3 +56,21 @@ def _get_extension_path(lib_name) -> str:
         raise ImportError
 
     return ext_specs.origin
+
+
+def _load_library(lib_name: str) -> None:
+    lib_path = _get_extension_path(lib_name)
+    # On Windows Python-3.8+ has `os.add_dll_directory` call,
+    # which is called from _get_extension_path to configure dll search path
+    # Condition below adds a workaround for older versions by
+    # explicitly calling `LoadLibraryExW` with the following flags:
+    #  - LOAD_LIBRARY_SEARCH_DEFAULT_DIRS (0x1000)
+    #  - LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR (0x100)
+    if os.name == "nt" and sys.version_info < (3, 8):
+
+        _kernel32 = ctypes.WinDLL("kernel32.dll", use_last_error=True)
+        if hasattr(_kernel32, "LoadLibraryExW"):
+            _image_lib_handle = _kernel32.LoadLibraryExW(lib_path, None, 0x00001100)
+        else:
+            warn("LoadLibraryExW is missing in kernel32.dll")
+    torch.ops.load_library(lib_path)
diff --git a/torchvision/io/_video_opt.py b/torchvision/io/_video_opt.py
@@ -5,12 +5,11 @@
 
 import torch
 
-from .._internally_replaced_utils import _get_extension_path
+from .._internally_replaced_utils import _load_library
 
 
 try:
-    lib_path = _get_extension_path("video_reader")
-    torch.ops.load_library(lib_path)
+    _load_library("video_reader")
     _HAS_VIDEO_OPT = True
 except (ImportError, OSError):
     _HAS_VIDEO_OPT = False

diff --git a/torchvision/io/image.py b/torchvision/io/image.py
@@ -6,25 +6,11 @@
 
 import torch
 
-from .._internally_replaced_utils import _get_extension_path
+from .._internally_replaced_utils import _load_library
 
 
 try:
-    lib_path = _get_extension_path("image")
-    # On Windows Python-3.8+ has `os.add_dll_directory` call,
-    # which is called from _get_extension_path to configure dll search path
-    # Condition below adds a workaround for older versions by
-    # explicitly calling `LoadLibraryExW` with the following flags:
-    #  - LOAD_LIBRARY_SEARCH_DEFAULT_DIRS (0x1000)
-    #  - LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR (0x100)
-    if os.name == "nt" and sys.version_info < (3, 8):
-
-        _kernel32 = ctypes.WinDLL("kernel32.dll", use_last_error=True)
-        if hasattr(_kernel32, "LoadLibraryExW"):
-            _image_lib_handle = _kernel32.LoadLibraryExW(lib_path, None, 0x00001100)
-        else:
-            warn("LoadLibraryExW is missing in kernel32.dll")
-    torch.ops.load_library(lib_path)
+    _load_library("image")
 except (ImportError, OSError) as e:
     warn(f"Failed to load {globals().get('lib_path','image.pyd')}: {e}")