Merge pull request #2420 from msabramo/wheel_check_required_packages

Refactor wheel command check of required packages

Author: Ivoz

pip/commands/wheel.py

@@ -9,7 +9,7 @@
 from pip.index import PackageFinder
 from pip.exceptions import CommandError, PreviousBuildDirError
 from pip.req import InstallRequirement, RequirementSet, parse_requirements
-from pip.utils import normalize_path
+from pip.utils import import_or_raise, normalize_path
 from pip.utils.build import BuildDirectory
 from pip.utils.deprecation import RemovedInPip7Warning, RemovedInPip8Warning
 from pip.wheel import WheelBuilder
@@ -100,33 +100,28 @@ def __init__(self, *args, **kw):
         self.parser.insert_option_group(0, index_opts)
         self.parser.insert_option_group(0, cmd_opts)
 
-    def run(self, options, args):
-
-        # confirm requirements
-        try:
-            import wheel.bdist_wheel
-            # Hack to make flake8 not complain about an unused import
-            wheel.bdist_wheel
-        except ImportError:
+    def check_required_packages(self):
+        import_or_raise(
+            'wheel.bdist_wheel',
+            CommandError,
+            "'pip wheel' requires the 'wheel' package. To fix this, run: "
+            "pip install wheel"
+        )
+        pkg_resources = import_or_raise(
+            'pkg_resources',
+            CommandError,
+            "'pip wheel' requires setuptools >= 0.8 for dist-info support."
+            " To fix this, run: pip install --upgrade setuptools"
+        )
+        if not hasattr(pkg_resources, 'DistInfoDistribution'):
             raise CommandError(
-                "'pip wheel' requires the 'wheel' package. To fix this, run: "
-                "pip install wheel"
+                "'pip wheel' requires setuptools >= 0.8 for dist-info "
+                "support. To fix this, run: pip install --upgrade "
+                "setuptools"
             )
 
-        try:
-            import pkg_resources
-        except ImportError:
-            raise CommandError(
-                "'pip wheel' requires setuptools >= 0.8 for dist-info support."
-                " To fix this, run: pip install --upgrade setuptools"
-            )
-        else:
-            if not hasattr(pkg_resources, 'DistInfoDistribution'):
-                raise CommandError(
-                    "'pip wheel' requires setuptools >= 0.8 for dist-info "
-                    "support. To fix this, run: pip install --upgrade "
-                    "setuptools"
-                )
+    def run(self, options, args):
+        self.check_required_packages()
 
         index_urls = [options.index_url] + options.extra_index_urls
         if options.no_index:

pip/utils/__init__.py

@@ -45,6 +45,13 @@
 logger = logging.getLogger(__name__)
 
 
+def import_or_raise(pkg_or_module_string, ExceptionType, *args, **kwargs):
+    try:
+        return __import__(pkg_or_module_string)
+    except ImportError:
+        raise ExceptionType(*args, **kwargs)
+
+
 def get_prog():
     try:
         if os.path.basename(sys.argv[0]) in ('__main__.py', '-c'):