[Needs discussion] Improve startup time through lazy-loading

shub/image/__init__.py

@@ -1,5 +1,7 @@
-import click
 import importlib
+import sys
+
+import click
 
 
 @click.group(help="Manage project based on custom Docker image")
@@ -18,6 +20,9 @@ def cli():
     "check",
 ]
 
+if len(sys.argv) > 2 and sys.argv[2] in module_deps:
+    module_deps = [sys.argv[2]]
+
 for command in module_deps:
     module_path = "shub.image." + command
     command_module = importlib.import_module(module_path)

shub/tool.py

@@ -1,5 +1,6 @@
 from __future__ import absolute_import
 import importlib
+import sys
 
 import click
 
@@ -51,6 +52,12 @@ def cli():
     "image",
 ]
 
+# Some imports, particularly requests and pip, are very slow. To avoid
+# importing these modules when running a command that doesn't need them, we
+# import that command module only.
+if len(sys.argv) > 1 and sys.argv[1] in commands:
+    commands = [sys.argv[1]]
+
 for command in commands:
     module_path = "shub." + command
     command_module = importlib.import_module(module_path)

shub/utils.py

@@ -19,18 +19,8 @@
 from six.moves.urllib.parse import urljoin
 
 import click
-import pip
-import requests
 import yaml
 
-from scrapinghub import Connection, APIError
-
-try:
-    from scrapinghub import HubstorageClient
-except ImportError:
-    # scrapinghub < 1.9.0
-    from hubstorage import HubstorageClient
-
 import shub
 from shub.compat import to_native_str
 from shub.exceptions import (BadParameterException, InvalidAuthException,
@@ -78,6 +68,7 @@ def create_default_setup_py(**kwargs):
 
 
 def make_deploy_request(url, data, files, auth, verbose, keep_log):
+    import requests
     last_logs = deque(maxlen=LAST_N_LOGS)
     try:
         rsp = requests.post(url=url, auth=auth, data=data, files=files,
@@ -284,6 +275,8 @@ def run_python(cmd, *args, **kwargs):
 
 
 def decompress_egg_files(directory=None):
+    # Import takes about half a second, so we do it lazily
+    import pip
     try:
         EXTS = pip.utils.ARCHIVE_EXTENSIONS
     except AttributeError:
@@ -399,6 +392,11 @@ def get_job_specs(job):
 
 
 def get_job(job):
+    try:
+        from scrapinghub import HubstorageClient
+    except ImportError:
+        # scrapinghub < 1.9.0
+        from hubstorage import HubstorageClient
     jobid, apikey = get_job_specs(job)
     hsc = HubstorageClient(auth=apikey)
     job = hsc.get_job(jobid)
@@ -553,6 +551,7 @@ def latest_github_release(force_update=False, timeout=1., cache=None):
         # saved
         if release_data.get('_shub_last_update', 0) == today:
             return release_data
+    import requests
     release_data = requests.get(REQ_URL, timeout=timeout).json()
     release_data['_shub_last_update'] = today
     try:
@@ -590,6 +589,8 @@ def update_available(silent_fail=True):
 
 
 def download_from_pypi(dest, pkg=None, reqfile=None, extra_args=None):
+    # Import takes about half a second, so we do it lazily
+    import pip
     if (not pkg and not reqfile) or (pkg and reqfile):
         raise ValueError('Call with either pkg or reqfile')
     extra_args = extra_args or []
@@ -643,6 +644,7 @@ def has_project_access(project, endpoint, apikey):
     """Check whether an API key has access to a given project. May raise
     InvalidAuthException if the API key is invalid (but not if it is valid but
     lacks access to the project)"""
+    from scrapinghub import Connection, APIError
     conn = Connection(apikey, url=endpoint)
     try:
         return project in conn.project_ids()

tests/test_utils.py

@@ -26,7 +26,7 @@
     NotFoundException, RemoteErrorException, SubcommandException
 )
 
-from .utils import AssertInvokeRaisesMixin, mock_conf
+from .utils import AssertInvokeRaisesMixin, mock_conf, mock_lazy_import
 
 
 class UtilsTest(AssertInvokeRaisesMixin, unittest.TestCase):
@@ -148,7 +148,7 @@ def test_get_job_specs_validates_jobid(self):
             with self.assertRaises(BadParameterException):
                 utils.get_job_specs(job_id)
 
-    @patch('shub.utils.HubstorageClient', autospec=True)
+    @patch('scrapinghub.HubstorageClient', autospec=True)
     def test_get_job(self, mock_HSC):
         class MockJob(object):
             metadata = {'some': 'val'}
@@ -271,7 +271,7 @@ def jri_result(follow, tail=None):
         job.resource.stats.return_value = {'totals': {'input_values': 1000}}
         self.assertEqual(jri_result(True, tail=3), [])
 
-    @patch('shub.utils.requests.get', autospec=True)
+    @patch('requests.get', autospec=True)
     def test_latest_github_release(self, mock_get):
         with self.runner.isolated_filesystem():
             mock_get.return_value.json.return_value = {'key': 'value'}
@@ -343,7 +343,7 @@ class MockException(Exception):
         with self.assertRaises(MockException):
             utils.update_available(silent_fail=False)
 
-    @patch('shub.utils.pip', autospec=True)
+    @mock_lazy_import('pip', autospec=True)
     def test_download_from_pypi(self, mock_pip):
         def _call(*args, **kwargs):
             utils.download_from_pypi(*args, **kwargs)
@@ -463,7 +463,7 @@ def call_update_yaml_dict():
             result = runner.invoke(call_update_yaml_dict)
         assert 'deprecated' in result.output
 
-    @patch('shub.utils.Connection')
+    @patch('scrapinghub.Connection')
     def test_has_project_access(self, mock_conn):
         mock_conn.return_value.project_ids.side_effect = APIError(
             'Authentication failed')

tests/utils.py

@@ -1,10 +1,12 @@
 # -*- coding: utf-8 -*-
 from __future__ import absolute_import
+import importlib
 import sys
 import re
 
 import mock
 from click.testing import CliRunner
+from mock import create_autospec, Mock, patch
 from tqdm._utils import _supports_unicode
 
 from shub import config
@@ -84,3 +86,17 @@ def clean_progress_output(output):
                       #    ("ESC"  + single command character)
         """,
         '', output)
+
+
+def mock_lazy_import(modname, autospec=False):
+    if autospec:
+        mod_mock = create_autospec(importlib.import_module(modname))
+    else:
+        mod_mock = Mock()
+
+    def decorator(function):
+        def wrapper(self, *args, **kwargs):
+            with patch.dict(sys.modules, {modname: mod_mock}):
+                return function(self, mod_mock, *args, **kwargs)
+        return wrapper
+    return decorator