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