replace imp with importlib (fixes #1607) (#1976)

Co-authored-by: Francisco Massa <fvsmassa@gmail.com>
pytorch · Mar 17, 2020 · 7a36388 · 7a36388
1 parent 32e1680
commit 7a36388
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 7 deletions.
diff --git a/torchvision/extension.py b/torchvision/extension.py
@@ -3,13 +3,21 @@
 
 def _register_extensions():
     import os
-    import imp
+    import importlib
     import torch
 
     # load the custom_op_library and register the custom ops
     lib_dir = os.path.dirname(__file__)
-    _, path, _ = imp.find_module("_C", [lib_dir])
-    torch.ops.load_library(path)
+    loader_details = (
+        importlib.machinery.ExtensionFileLoader,
+        importlib.machinery.EXTENSION_SUFFIXES
+    )
+
+    extfinder = importlib.machinery.FileFinder(lib_dir, loader_details)
+    ext_specs = extfinder.find_spec("_C")
+    if ext_specs is None:
+        raise ImportError
+    torch.ops.load_library(ext_specs.origin)
 
 
 try:

diff --git a/torchvision/io/_video_opt.py b/torchvision/io/_video_opt.py
@@ -1,5 +1,5 @@
 
-import imp
+import importlib
 import math
 import os
 import warnings
@@ -14,9 +14,17 @@
 
 try:
     lib_dir = os.path.join(os.path.dirname(__file__), "..")
-    _, path, description = imp.find_module("video_reader", [lib_dir])
-    torch.ops.load_library(path)
-    _HAS_VIDEO_OPT = True
+
+    loader_details = (
+        importlib.machinery.ExtensionFileLoader,
+        importlib.machinery.EXTENSION_SUFFIXES
+    )
+
+    extfinder = importlib.machinery.FileFinder(lib_dir, loader_details)
+    ext_specs = extfinder.find_spec("video_reader")
+    if ext_specs is not None:
+        torch.ops.load_library(ext_specs.origin)
+        _HAS_VIDEO_OPT = True
 except (ImportError, OSError):
     pass