Fix pip environment
[yardstick.git] / ez_setup.py
1 #!/usr/bin/env python
2 """Bootstrap setuptools installation
3
4 To use setuptools in your package's setup.py, include this
5 file in the same directory and add this to the top of your setup.py::
6
7     from ez_setup import use_setuptools
8     use_setuptools()
9
10 To require a specific version of setuptools, set a download
11 mirror, or use an alternate download directory, simply supply
12 the appropriate options to ``use_setuptools()``.
13
14 This file can also be run as a script to install or upgrade setuptools.
15 """
16 from __future__ import absolute_import
17 import os
18 import shutil
19 import sys
20 import tempfile
21 import zipfile
22 import optparse
23 import subprocess
24 import platform
25 import contextlib
26
27 from distutils import log
28
29 try:
30     from urllib.request import urlopen
31 except ImportError:
32     from six.moves.urllib import urlopen
33
34 try:
35     from site import USER_SITE
36 except ImportError:
37     USER_SITE = None
38
39 DEFAULT_VERSION = "6.1"
40 DEFAULT_URL = "https://pypi.python.org/packages/source/s/setuptools/"
41
42
43 def _python_cmd(*args):
44     """
45     Return True if the command succeeded.
46     """
47     args = (sys.executable,) + args
48     return subprocess.call(args) == 0
49
50
51 def _install(archive_filename, install_args=()):
52     with archive_context(archive_filename):
53         # installing
54         log.warn('Installing Setuptools')
55         if not _python_cmd('setup.py', 'install', *install_args):
56             log.warn('Something went wrong during the installation.')
57             log.warn('See the error message above.')
58             # exitcode will be 2
59             return 2
60
61
62 def _build_egg(egg, archive_filename, to_dir):
63     with archive_context(archive_filename):
64         # building an egg
65         log.warn('Building a Setuptools egg in %s', to_dir)
66         _python_cmd('setup.py', '-q', 'bdist_egg', '--dist-dir', to_dir)
67     # returning the result
68     log.warn(egg)
69     if not os.path.exists(egg):
70         raise IOError('Could not build the egg.')
71
72
73 class ContextualZipFile(zipfile.ZipFile):
74     """
75     Supplement ZipFile class to support context manager for Python 2.6
76     """
77
78     def __enter__(self):
79         return self
80
81     def __exit__(self, type, value, traceback):
82         self.close()
83
84     def __new__(cls, *args, **kwargs):
85         """
86         Construct a ZipFile or ContextualZipFile as appropriate
87         """
88         if hasattr(zipfile.ZipFile, '__exit__'):
89             return zipfile.ZipFile(*args, **kwargs)
90         return super(ContextualZipFile, cls).__new__(cls)
91
92
93 @contextlib.contextmanager
94 def archive_context(filename):
95     # extracting the archive
96     tmpdir = tempfile.mkdtemp()
97     log.warn('Extracting in %s', tmpdir)
98     old_wd = os.getcwd()
99     try:
100         os.chdir(tmpdir)
101         with ContextualZipFile(filename) as archive:
102             archive.extractall()
103
104         # going in the directory
105         subdir = os.path.join(tmpdir, os.listdir(tmpdir)[0])
106         os.chdir(subdir)
107         log.warn('Now working in %s', subdir)
108         yield
109
110     finally:
111         os.chdir(old_wd)
112         shutil.rmtree(tmpdir)
113
114
115 def _do_download(version, download_base, to_dir, download_delay):
116     egg = os.path.join(to_dir, 'setuptools-%s-py%d.%d.egg'
117                        % (version, sys.version_info[0], sys.version_info[1]))
118     if not os.path.exists(egg):
119         archive = download_setuptools(version, download_base,
120                                       to_dir, download_delay)
121         _build_egg(egg, archive, to_dir)
122     sys.path.insert(0, egg)
123
124     # Remove previously-imported pkg_resources if present (see
125     # https://bitbucket.org/pypa/setuptools/pull-request/7/ for details).
126     if 'pkg_resources' in sys.modules:
127         del sys.modules['pkg_resources']
128
129     import setuptools
130     setuptools.bootstrap_install_from = egg
131
132
133 def use_setuptools(version=DEFAULT_VERSION, download_base=DEFAULT_URL,
134                    to_dir=os.curdir, download_delay=15):
135     to_dir = os.path.abspath(to_dir)
136     rep_modules = 'pkg_resources', 'setuptools'
137     imported = set(sys.modules).intersection(rep_modules)
138     try:
139         import pkg_resources
140     except ImportError:
141         return _do_download(version, download_base, to_dir, download_delay)
142     try:
143         pkg_resources.require("setuptools>=" + version)
144         return
145     except pkg_resources.DistributionNotFound:
146         return _do_download(version, download_base, to_dir, download_delay)
147     except pkg_resources.VersionConflict as VC_err:
148         if imported:
149             msg = """\
150 The required version of setuptools (>={version}) is not available,
151 and can't be installed while this script is running. Please
152 install a more recent version first, using
153 'easy_install -U setuptools'.
154
155 (Currently using {VC_err.args[0]!r})
156 """.format(VC_err=VC_err, version=version)
157             sys.stderr.write(msg)
158             sys.exit(2)
159
160         # otherwise, reload ok
161         del pkg_resources, sys.modules['pkg_resources']
162         return _do_download(version, download_base, to_dir, download_delay)
163
164
165 def _clean_check(cmd, target):
166     """
167     Run the command to download target. If the command fails, clean up before
168     re-raising the error.
169     """
170     try:
171         subprocess.check_call(cmd)
172     except subprocess.CalledProcessError:
173         if os.access(target, os.F_OK):
174             os.unlink(target)
175         raise
176
177
178 def download_file_powershell(url, target):
179     """
180     Download the file at url to target using Powershell (which will validate
181     trust). Raise an exception if the command cannot complete.
182     """
183     target = os.path.abspath(target)
184     ps_cmd = (
185         "[System.Net.WebRequest]::DefaultWebProxy.Credentials = "
186         "[System.Net.CredentialCache]::DefaultCredentials; "
187         "(new-object System.Net.WebClient).DownloadFile(%(url)r, %(target)r)"
188         % vars()
189     )
190     cmd = [
191         'powershell',
192         '-Command',
193         ps_cmd,
194     ]
195     _clean_check(cmd, target)
196
197
198 def has_powershell():
199     if platform.system() != 'Windows':
200         return False
201     cmd = ['powershell', '-Command', 'echo test']
202     with open(os.path.devnull, 'wb') as devnull:
203         try:
204             subprocess.check_call(cmd, stdout=devnull, stderr=devnull)
205         except Exception:
206             return False
207     return True
208
209
210 download_file_powershell.viable = has_powershell
211
212
213 def download_file_curl(url, target):
214     cmd = ['curl', url, '--silent', '--output', target]
215     _clean_check(cmd, target)
216
217
218 def has_curl():
219     cmd = ['curl', '--version']
220     with open(os.path.devnull, 'wb') as devnull:
221         try:
222             subprocess.check_call(cmd, stdout=devnull, stderr=devnull)
223         except Exception:
224             return False
225     return True
226
227
228 download_file_curl.viable = has_curl
229
230
231 def download_file_wget(url, target):
232     cmd = ['wget', url, '--quiet', '--output-document', target]
233     _clean_check(cmd, target)
234
235
236 def has_wget():
237     cmd = ['wget', '--version']
238     with open(os.path.devnull, 'wb') as devnull:
239         try:
240             subprocess.check_call(cmd, stdout=devnull, stderr=devnull)
241         except Exception:
242             return False
243     return True
244
245
246 download_file_wget.viable = has_wget
247
248
249 def download_file_insecure(url, target):
250     """
251     Use Python to download the file, even though it cannot authenticate the
252     connection.
253     """
254     src = urlopen(url)
255     try:
256         # Read all the data in one block.
257         data = src.read()
258     finally:
259         src.close()
260
261     # Write all the data in one block to avoid creating a partial file.
262     with open(target, "wb") as dst:
263         dst.write(data)
264
265
266 download_file_insecure.viable = lambda: True
267
268
269 def get_best_downloader():
270     downloaders = (
271         download_file_powershell,
272         download_file_curl,
273         download_file_wget,
274         download_file_insecure,
275     )
276     viable_downloaders = (dl for dl in downloaders if dl.viable())
277     return next(viable_downloaders, None)
278
279
280 def download_setuptools(version=DEFAULT_VERSION, download_base=DEFAULT_URL,
281                         to_dir=os.curdir, delay=15,
282                         downloader_factory=get_best_downloader):
283     """
284     Download setuptools from a specified location and return its filename
285
286     `version` should be a valid setuptools version number that is available
287     as an sdist for download under the `download_base` URL (which should end
288     with a '/'). `to_dir` is the directory where the egg will be downloaded.
289     `delay` is the number of seconds to pause before an actual download
290     attempt.
291
292     ``downloader_factory`` should be a function taking no arguments and
293     returning a function for downloading a URL to a target.
294     """
295     # making sure we use the absolute path
296     to_dir = os.path.abspath(to_dir)
297     zip_name = "setuptools-%s.zip" % version
298     url = download_base + zip_name
299     saveto = os.path.join(to_dir, zip_name)
300     if not os.path.exists(saveto):  # Avoid repeated downloads
301         log.warn("Downloading %s", url)
302         downloader = downloader_factory()
303         downloader(url, saveto)
304     return os.path.realpath(saveto)
305
306
307 def _build_install_args(options):
308     """
309     Build the arguments to 'python setup.py install' on the setuptools package
310     """
311     return ['--user'] if options.user_install else []
312
313
314 def _parse_args():
315     """
316     Parse the command line for options
317     """
318     parser = optparse.OptionParser()
319     parser.add_option(
320         '--user', dest='user_install', action='store_true', default=False,
321         help='install in user site package (requires Python 2.6 or later)')
322     parser.add_option(
323         '--download-base', dest='download_base', metavar="URL",
324         default=DEFAULT_URL,
325         help='alternative URL from where to download the setuptools package')
326     parser.add_option(
327         '--insecure', dest='downloader_factory', action='store_const',
328         const=lambda: download_file_insecure, default=get_best_downloader,
329         help='Use internal, non-validating downloader'
330     )
331     parser.add_option(
332         '--version', help="Specify which version to download",
333         default=DEFAULT_VERSION,
334     )
335     options, args = parser.parse_args()
336     # positional arguments are ignored
337     return options
338
339
340 def main():
341     """Install or upgrade setuptools and EasyInstall"""
342     options = _parse_args()
343     archive = download_setuptools(
344         version=options.version,
345         download_base=options.download_base,
346         downloader_factory=options.downloader_factory,
347     )
348     return _install(archive, _build_install_args(options))
349
350
351 if __name__ == '__main__':
352     sys.exit(main())