1#!/usr/bin/env python
2
3"""
4Setuptools bootstrapping installer.
5
6Run this script to install or upgrade setuptools.
7"""
8
9import os
10import shutil
11import sys
12import tempfile
13import zipfile
14import optparse
15import subprocess
16import platform
17import textwrap
18import contextlib
19import warnings
20
21from distutils import log
22
23try:
24    from urllib.request import urlopen
25except ImportError:
26    from urllib2 import urlopen
27
28try:
29    from site import USER_SITE
30except ImportError:
31    USER_SITE = None
32
33DEFAULT_VERSION = "14.3.1"
34DEFAULT_URL = "https://pypi.python.org/packages/source/s/setuptools/"
35DEFAULT_SAVE_DIR = os.curdir
36
37
38def _python_cmd(*args):
39    """
40    Execute a command.
41
42    Return True if the command succeeded.
43    """
44    args = (sys.executable,) + args
45    return subprocess.call(args) == 0
46
47
48def _install(archive_filename, install_args=()):
49    """Install Setuptools."""
50    with archive_context(archive_filename):
51        # installing
52        log.warn("Installing Setuptools")
53        if not _python_cmd("setup.py", "install", *install_args):
54            log.warn("Something went wrong during the installation.")
55            log.warn("See the error message above.")
56            # exitcode will be 2
57            return 2
58
59
60def _build_egg(egg, archive_filename, to_dir):
61    """Build Setuptools egg."""
62    with archive_context(archive_filename):
63        # building an egg
64        log.warn("Building a Setuptools egg in %s", to_dir)
65        _python_cmd("setup.py", "-q", "bdist_egg", "--dist-dir", to_dir)
66    # returning the result
67    log.warn(egg)
68    if not os.path.exists(egg):
69        raise IOError("Could not build the egg.")
70
71
72class ContextualZipFile(zipfile.ZipFile):
73
74    """Supplement ZipFile class to support context manager for Python 2.6."""
75
76    def __enter__(self):
77        return self
78
79    def __exit__(self, type, value, traceback):
80        self.close()
81
82    def __new__(cls, *args, **kwargs):
83        """Construct a ZipFile or ContextualZipFile as appropriate."""
84        if hasattr(zipfile.ZipFile, "__exit__"):
85            return zipfile.ZipFile(*args, **kwargs)
86        return super(ContextualZipFile, cls).__new__(cls)
87
88
89@contextlib.contextmanager
90def archive_context(filename):
91    """
92    Unzip filename to a temporary directory, set to the cwd.
93
94    The unzipped target is cleaned up after.
95    """
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
115def _do_download(version, download_base, to_dir, download_delay):
116    """Download Setuptools."""
117    egg = os.path.join(
118        to_dir,
119        "setuptools-%s-py%d.%d.egg"
120        % (version, sys.version_info[0], sys.version_info[1]),
121    )
122    if not os.path.exists(egg):
123        archive = download_setuptools(version, download_base, to_dir, download_delay)
124        _build_egg(egg, archive, to_dir)
125    sys.path.insert(0, egg)
126
127    # Remove previously-imported pkg_resources if present (see
128    # https://bitbucket.org/pypa/setuptools/pull-request/7/ for details).
129    if "pkg_resources" in sys.modules:
130        del sys.modules["pkg_resources"]
131
132    import setuptools
133
134    setuptools.bootstrap_install_from = egg
135
136
137def use_setuptools(
138    version=DEFAULT_VERSION,
139    download_base=DEFAULT_URL,
140    to_dir=DEFAULT_SAVE_DIR,
141    download_delay=15,
142):
143    """
144    Ensure that a setuptools version is installed.
145
146    Return None. Raise SystemExit if the requested version
147    or later cannot be installed.
148    """
149    to_dir = os.path.abspath(to_dir)
150
151    # prior to importing, capture the module state for
152    # representative modules.
153    rep_modules = "pkg_resources", "setuptools"
154    imported = set(sys.modules).intersection(rep_modules)
155
156    try:
157        import pkg_resources
158
159        pkg_resources.require("setuptools>=" + version)
160        # a suitable version is already installed
161        return
162    except ImportError:
163        # pkg_resources not available; setuptools is not installed; download
164        pass
165    except pkg_resources.DistributionNotFound:
166        # no version of setuptools was found; allow download
167        pass
168    except pkg_resources.VersionConflict as VC_err:
169        if imported:
170            _conflict_bail(VC_err, version)
171
172        # otherwise, unload pkg_resources to allow the downloaded version to
173        #  take precedence.
174        del pkg_resources
175        _unload_pkg_resources()
176
177    return _do_download(version, download_base, to_dir, download_delay)
178
179
180def _conflict_bail(VC_err, version):
181    """
182    Setuptools was imported prior to invocation, so it is
183    unsafe to unload it. Bail out.
184    """
185    conflict_tmpl = textwrap.dedent(
186        """
187        The required version of setuptools (>={version}) is not available,
188        and can't be installed while this script is running. Please
189        install a more recent version first, using
190        'easy_install -U setuptools'.
191
192        (Currently using {VC_err.args[0]!r})
193        """
194    )
195    msg = conflict_tmpl.format(**locals())
196    sys.stderr.write(msg)
197    sys.exit(2)
198
199
200def _unload_pkg_resources():
201    del_modules = [name for name in sys.modules if name.startswith("pkg_resources")]
202    for mod_name in del_modules:
203        del sys.modules[mod_name]
204
205
206def _clean_check(cmd, target):
207    """
208    Run the command to download target.
209
210    If the command fails, clean up before re-raising the error.
211    """
212    try:
213        subprocess.check_call(cmd)
214    except subprocess.CalledProcessError:
215        if os.access(target, os.F_OK):
216            os.unlink(target)
217        raise
218
219
220def download_file_powershell(url, target):
221    """
222    Download the file at url to target using Powershell.
223
224    Powershell will validate trust.
225    Raise an exception if the command cannot complete.
226    """
227    target = os.path.abspath(target)
228    ps_cmd = (
229        "[System.Net.WebRequest]::DefaultWebProxy.Credentials = "
230        "[System.Net.CredentialCache]::DefaultCredentials; "
231        "(new-object System.Net.WebClient).DownloadFile(%(url)r, %(target)r)" % vars()
232    )
233    cmd = ["powershell", "-Command", ps_cmd]
234    _clean_check(cmd, target)
235
236
237def has_powershell():
238    """Determine if Powershell is available."""
239    if platform.system() != "Windows":
240        return False
241    cmd = ["powershell", "-Command", "echo test"]
242    with open(os.path.devnull, "wb") as devnull:
243        try:
244            subprocess.check_call(cmd, stdout=devnull, stderr=devnull)
245        except Exception:
246            return False
247    return True
248
249
250download_file_powershell.viable = has_powershell
251
252
253def download_file_curl(url, target):
254    cmd = ["curl", url, "--silent", "--output", target]
255    _clean_check(cmd, target)
256
257
258def has_curl():
259    cmd = ["curl", "--version"]
260    with open(os.path.devnull, "wb") as devnull:
261        try:
262            subprocess.check_call(cmd, stdout=devnull, stderr=devnull)
263        except Exception:
264            return False
265    return True
266
267
268download_file_curl.viable = has_curl
269
270
271def download_file_wget(url, target):
272    cmd = ["wget", url, "--quiet", "--output-document", target]
273    _clean_check(cmd, target)
274
275
276def has_wget():
277    cmd = ["wget", "--version"]
278    with open(os.path.devnull, "wb") as devnull:
279        try:
280            subprocess.check_call(cmd, stdout=devnull, stderr=devnull)
281        except Exception:
282            return False
283    return True
284
285
286download_file_wget.viable = has_wget
287
288
289def download_file_insecure(url, target):
290    """Use Python to download the file, without connection authentication."""
291    src = urlopen(url)
292    try:
293        # Read all the data in one block.
294        data = src.read()
295    finally:
296        src.close()
297
298    # Write all the data in one block to avoid creating a partial file.
299    with open(target, "wb") as dst:
300        dst.write(data)
301
302
303download_file_insecure.viable = lambda: True
304
305
306def get_best_downloader():
307    downloaders = (
308        download_file_powershell,
309        download_file_curl,
310        download_file_wget,
311        download_file_insecure,
312    )
313    viable_downloaders = (dl for dl in downloaders if dl.viable())
314    return next(viable_downloaders, None)
315
316
317def download_setuptools(
318    version=DEFAULT_VERSION,
319    download_base=DEFAULT_URL,
320    to_dir=DEFAULT_SAVE_DIR,
321    delay=15,
322    downloader_factory=get_best_downloader,
323):
324    """
325    Download setuptools from a specified location and return its filename.
326
327    `version` should be a valid setuptools version number that is available
328    as an sdist for download under the `download_base` URL (which should end
329    with a '/'). `to_dir` is the directory where the egg will be downloaded.
330    `delay` is the number of seconds to pause before an actual download
331    attempt.
332
333    ``downloader_factory`` should be a function taking no arguments and
334    returning a function for downloading a URL to a target.
335    """
336    # making sure we use the absolute path
337    to_dir = os.path.abspath(to_dir)
338    zip_name = "setuptools-%s.zip" % version
339    url = download_base + zip_name
340    saveto = os.path.join(to_dir, zip_name)
341    if not os.path.exists(saveto):  # Avoid repeated downloads
342        log.warn("Downloading %s", url)
343        downloader = downloader_factory()
344        downloader(url, saveto)
345    return os.path.realpath(saveto)
346
347
348def _build_install_args(options):
349    """
350    Build the arguments to 'python setup.py install' on the setuptools package.
351
352    Returns list of command line arguments.
353    """
354    return ["--user"] if options.user_install else []
355
356
357def _parse_args():
358    """Parse the command line for options."""
359    parser = optparse.OptionParser()
360    parser.add_option(
361        "--user",
362        dest="user_install",
363        action="store_true",
364        default=False,
365        help="install in user site package (requires Python 2.6 or later)",
366    )
367    parser.add_option(
368        "--download-base",
369        dest="download_base",
370        metavar="URL",
371        default=DEFAULT_URL,
372        help="alternative URL from where to download the setuptools package",
373    )
374    parser.add_option(
375        "--insecure",
376        dest="downloader_factory",
377        action="store_const",
378        const=lambda: download_file_insecure,
379        default=get_best_downloader,
380        help="Use internal, non-validating downloader",
381    )
382    parser.add_option(
383        "--version", help="Specify which version to download", default=DEFAULT_VERSION
384    )
385    parser.add_option(
386        "--to-dir",
387        help="Directory to save (and re-use) package",
388        default=DEFAULT_SAVE_DIR,
389    )
390    options, args = parser.parse_args()
391    # positional arguments are ignored
392    return options
393
394
395def _download_args(options):
396    """Return args for download_setuptools function from cmdline args."""
397    return dict(
398        version=options.version,
399        download_base=options.download_base,
400        downloader_factory=options.downloader_factory,
401        to_dir=options.to_dir,
402    )
403
404
405def main():
406    """Install or upgrade setuptools and EasyInstall."""
407    options = _parse_args()
408    archive = download_setuptools(**_download_args(options))
409    return _install(archive, _build_install_args(options))
410
411
412if __name__ == "__main__":
413    sys.exit(main())
414