Extract 'imp' re-implementation to setuptools._imp and wrap it in py27compat for compatibility.

4 files changed, 113 insertions, 87 deletions

diff --git a/conftest.py b/conftest.py
index 0d7b274c..1746bfb5 100644
--- a/conftest.py
+++ b/conftest.py
@@ -19,6 +19,7 @@ collect_ignore = [
 
 if sys.version_info < (3,):
     collect_ignore.append('setuptools/lib2to3_ex.py')
+    collect_ignore.append('setuptools/_imp.py')
 
 
 if sys.version_info < (3, 6):
diff --git a/setuptools/_imp.py b/setuptools/_imp.py
new file mode 100644
index 00000000..6bc90243
--- /dev/null
+++ b/setuptools/_imp.py
@@ -0,0 +1,72 @@
+"""
+Re-implementation of find_module and get_frozen_object
+from the deprecated imp module.
+"""
+
+import os
+import importlib.util
+import importlib.machinery
+
+
+PY_SOURCE = 1
+PY_COMPILED = 2
+C_EXTENSION = 3
+C_BUILTIN = 6
+PY_FROZEN = 7
+
+
+def find_module(module, paths=None):
+    """
+    """
+    spec = importlib.util.find_spec(module, paths)
+    if spec is None:
+        raise ImportError("Can't find %s" % module)
+    if not spec.has_location and hasattr(spec, 'submodule_search_locations'):
+        spec = importlib.util.spec_from_loader('__init__.py', spec.loader)
+
+    kind = -1
+    file = None
+    static = isinstance(spec.loader, type)
+    if spec.origin == 'frozen' or static and issubclass(
+            spec.loader, importlib.machinery.FrozenImporter):
+        kind = PY_FROZEN
+        path = None  # imp compabilty
+        suffix = mode = ''  # imp compability
+    elif spec.origin == 'built-in' or static and issubclass(
+            spec.loader, importlib.machinery.BuiltinImporter):
+        kind = C_BUILTIN
+        path = None  # imp compabilty
+        suffix = mode = ''  # imp compability
+    elif spec.has_location:
+        path = spec.origin
+        suffix = os.path.splitext(path)[1]
+        mode = 'r' if suffix in importlib.machinery.SOURCE_SUFFIXES else 'rb'
+
+        if suffix in importlib.machinery.SOURCE_SUFFIXES:
+            kind = PY_SOURCE
+        elif suffix in importlib.machinery.BYTECODE_SUFFIXES:
+            kind = PY_COMPILED
+        elif suffix in importlib.machinery.EXTENSION_SUFFIXES:
+            kind = C_EXTENSION
+
+        if kind in {PY_SOURCE, PY_COMPILED}:
+            file = open(path, mode)
+    else:
+        path = None
+        suffix = mode = ''
+
+    return file, path, (suffix, mode, kind)
+
+
+def get_frozen_object(module, paths):
+    spec = importlib.util.find_spec(module, paths)
+    if hasattr(spec, 'submodule_search_locations'):
+        spec = importlib.util.spec_from_loader('__init__.py', spec.loader)
+    return spec.loader.get_code(module)
+
+
+def get_module(module, paths, info):
+    spec = importlib.util.find_spec(module, paths)
+    if hasattr(spec, 'submodule_search_locations'):
+        spec = importlib.util.spec_from_loader('__init__.py', spec.loader)
+    return importlib.util.module_from_spec(spec)
diff --git a/setuptools/depends.py b/setuptools/depends.py
index 97f0ed9d..eed4913a 100644
--- a/setuptools/depends.py
+++ b/setuptools/depends.py
@@ -1,22 +1,11 @@
 import sys
 import marshal
 from distutils.version import StrictVersion
-from setuptools.extern import six
 
 from .py33compat import Bytecode
 
-if six.PY2:
-    import imp
-    from imp import PKG_DIRECTORY, PY_COMPILED, PY_SOURCE, PY_FROZEN
-else:
-    import os.path
-    from importlib.util import find_spec, spec_from_loader
-    from importlib.machinery import SOURCE_SUFFIXES, BYTECODE_SUFFIXES, EXTENSION_SUFFIXES, BuiltinImporter, FrozenImporter
-    PY_SOURCE = 1
-    PY_COMPILED = 2
-    C_EXTENSION = 3
-    C_BUILTIN = 6
-    PY_FROZEN = 7
+from .py27compat import find_module, PY_COMPILED, PY_FROZEN, PY_SOURCE
+from . import py27compat
 
 
 __all__ = [
@@ -27,7 +16,8 @@ __all__ = [
 class Require:
     """A prerequisite to building or installing a distribution"""
 
-    def __init__(self, name, requested_version, module, homepage='',
+    def __init__(
+            self, name, requested_version, module, homepage='',
             attribute=None, format=None):
 
         if format is None and requested_version is not None:
@@ -91,63 +81,6 @@ class Require:
         return self.version_ok(version)
 
 
-def find_module(module, paths=None):
-    """Just like 'imp.find_module()', but with package support"""
-    if six.PY3:
-        spec = find_spec(module, paths)
-        if spec is None:
-            raise ImportError("Can't find %s" % module)
-        if not spec.has_location and hasattr(spec, 'submodule_search_locations'):
-            spec = spec_from_loader('__init__.py', spec.loader)
-
-        kind = -1
-        file = None
-        static = isinstance(spec.loader, type)
-        if spec.origin == 'frozen' or static and issubclass(spec.loader, FrozenImporter):
-            kind = PY_FROZEN
-            path = None # imp compabilty
-            suffix = mode = '' # imp compability
-        elif spec.origin == 'built-in' or static and issubclass(spec.loader, BuiltinImporter):
-            kind = C_BUILTIN
-            path = None # imp compabilty
-            suffix = mode = '' # imp compability
-        elif spec.has_location:
-            frozen = False
-            path = spec.origin
-            suffix = os.path.splitext(path)[1]
-            mode = 'r' if suffix in SOURCE_SUFFIXES else 'rb'
-
-            if suffix in SOURCE_SUFFIXES:
-                kind = PY_SOURCE
-            elif suffix in BYTECODE_SUFFIXES:
-                kind = PY_COMPILED
-            elif suffix in EXTENSION_SUFFIXES:
-                kind = C_EXTENSION
-
-            if kind in {PY_SOURCE, PY_COMPILED}:
-                file = open(path, mode)
-        else:
-            path = None
-            suffix = mode= ''
-
-        return file, path, (suffix, mode, kind)
-
-    else:
-        parts = module.split('.')
-        while parts:
-            part = parts.pop(0)
-            f, path, (suffix, mode, kind) = info = imp.find_module(part, paths)
-
-            if kind == PKG_DIRECTORY:
-                parts = parts or ['__init__']
-                paths = [path]
-
-            elif parts:
-                raise ImportError("Can't find %r in %s" % (parts, module))
-
-        return info
-
-
 def get_module_constant(module, symbol, default=-1, paths=None):
     """Find 'module' by searching 'paths', and extract 'symbol'
 
@@ -156,35 +89,23 @@ def get_module_constant(module, symbol, default=-1, paths=None):
     constant.  Otherwise, return 'default'."""
 
     try:
-        f, path, (suffix, mode, kind) = find_module(module, paths)
+        f, path, (suffix, mode, kind) = info = find_module(module, paths)
     except ImportError:
         # Module doesn't exist
         return None
 
-    if six.PY3:
-        spec = find_spec(module, paths)
-        if hasattr(spec, 'submodule_search_locations'):
-            spec = spec_from_loader('__init__.py', spec.loader)
-
     try:
         if kind == PY_COMPILED:
             f.read(8)  # skip magic & date
             code = marshal.load(f)
         elif kind == PY_FROZEN:
-            if six.PY2:
-                code = imp.get_frozen_object(module)
-            else:
-                code = spec.loader.get_code(module)
+            code = py27compat.get_frozen_object(module, paths)
         elif kind == PY_SOURCE:
             code = compile(f.read(), path, 'exec')
         else:
             # Not something we can parse; we'll have to import it.  :(
-            if module not in sys.modules:
-                if six.PY2:
-                    imp.load_module(module, f, path, (suffix, mode, kind))
-                else:
-                    sys.modules[module] = module_from_spec(spec)
-            return getattr(sys.modules[module], symbol, None)
+            imported = py27compat.get_module(module, paths, info)
+            return getattr(imported, symbol, None)
 
     finally:
         if f:
diff --git a/setuptools/py27compat.py b/setuptools/py27compat.py
index 2985011b..cf5fb33e 100644
--- a/setuptools/py27compat.py
+++ b/setuptools/py27compat.py
@@ -2,6 +2,7 @@
 Compatibility Support for Python 2.7 and earlier
 """
 
+import sys
 import platform
 
 from setuptools.extern import six
@@ -26,3 +27,34 @@ linux_py2_ascii = (
 
 rmtree_safe = str if linux_py2_ascii else lambda x: x
 """Workaround for http://bugs.python.org/issue24672"""
+
+
+try:
+    from ._imp import find_module, PY_COMPILED, PY_FROZEN, PY_SOURCE
+    from ._imp import get_frozen_object, get_module
+except ImportError:
+    import imp
+    from imp import PY_COMPILED, PY_FROZEN, PY_SOURCE  # noqa
+
+    def find_module(module, paths=None):
+        """Just like 'imp.find_module()', but with package support"""
+        parts = module.split('.')
+        while parts:
+            part = parts.pop(0)
+            f, path, (suffix, mode, kind) = info = imp.find_module(part, paths)
+
+            if kind == imp.PKG_DIRECTORY:
+                parts = parts or ['__init__']
+                paths = [path]
+
+            elif parts:
+                raise ImportError("Can't find %r in %s" % (parts, module))
+
+        return info
+
+    def get_frozen_object(module, paths):
+        return imp.get_frozen_object(module)
+
+    def get_module(module, paths, info):
+        imp.load_module(*info)
+        return sys.modules[module]