1"""
2Virtual environment (venv) package for Python. Based on PEP 405.
3
4Copyright (C) 2011-2014 Vinay Sajip.
5Licensed to the PSF under a contributor agreement.
6"""
7import logging
8import os
9import shutil
10import subprocess
11import sys
12import sysconfig
13import types
14
15
16CORE_VENV_DEPS = ('pip', 'setuptools')
17logger = logging.getLogger(__name__)
18
19
20class EnvBuilder:
21    """
22    This class exists to allow virtual environment creation to be
23    customized. The constructor parameters determine the builder's
24    behaviour when called upon to create a virtual environment.
25
26    By default, the builder makes the system (global) site-packages dir
27    *un*available to the created environment.
28
29    If invoked using the Python -m option, the default is to use copying
30    on Windows platforms but symlinks elsewhere. If instantiated some
31    other way, the default is to *not* use symlinks.
32
33    :param system_site_packages: If True, the system (global) site-packages
34                                 dir is available to created environments.
35    :param clear: If True, delete the contents of the environment directory if
36                  it already exists, before environment creation.
37    :param symlinks: If True, attempt to symlink rather than copy files into
38                     virtual environment.
39    :param upgrade: If True, upgrade an existing virtual environment.
40    :param with_pip: If True, ensure pip is installed in the virtual
41                     environment
42    :param prompt: Alternative terminal prefix for the environment.
43    :param upgrade_deps: Update the base venv modules to the latest on PyPI
44    """
45
46    def __init__(self, system_site_packages=False, clear=False,
47                 symlinks=False, upgrade=False, with_pip=False, prompt=None,
48                 upgrade_deps=False):
49        self.system_site_packages = system_site_packages
50        self.clear = clear
51        self.symlinks = symlinks
52        self.upgrade = upgrade
53        self.with_pip = with_pip
54        if prompt == '.':  # see bpo-38901
55            prompt = os.path.basename(os.getcwd())
56        self.prompt = prompt
57        self.upgrade_deps = upgrade_deps
58
59    def create(self, env_dir):
60        """
61        Create a virtual environment in a directory.
62
63        :param env_dir: The target directory to create an environment in.
64
65        """
66        env_dir = os.path.abspath(env_dir)
67        context = self.ensure_directories(env_dir)
68        # See issue 24875. We need system_site_packages to be False
69        # until after pip is installed.
70        true_system_site_packages = self.system_site_packages
71        self.system_site_packages = False
72        self.create_configuration(context)
73        self.setup_python(context)
74        if self.with_pip:
75            self._setup_pip(context)
76        if not self.upgrade:
77            self.setup_scripts(context)
78            self.post_setup(context)
79        if true_system_site_packages:
80            # We had set it to False before, now
81            # restore it and rewrite the configuration
82            self.system_site_packages = True
83            self.create_configuration(context)
84        if self.upgrade_deps:
85            self.upgrade_dependencies(context)
86
87    def clear_directory(self, path):
88        for fn in os.listdir(path):
89            fn = os.path.join(path, fn)
90            if os.path.islink(fn) or os.path.isfile(fn):
91                os.remove(fn)
92            elif os.path.isdir(fn):
93                shutil.rmtree(fn)
94
95    def ensure_directories(self, env_dir):
96        """
97        Create the directories for the environment.
98
99        Returns a context object which holds paths in the environment,
100        for use by subsequent logic.
101        """
102
103        def create_if_needed(d):
104            if not os.path.exists(d):
105                os.makedirs(d)
106            elif os.path.islink(d) or os.path.isfile(d):
107                raise ValueError('Unable to create directory %r' % d)
108
109        if os.path.exists(env_dir) and self.clear:
110            self.clear_directory(env_dir)
111        context = types.SimpleNamespace()
112        context.env_dir = env_dir
113        context.env_name = os.path.split(env_dir)[1]
114        prompt = self.prompt if self.prompt is not None else context.env_name
115        context.prompt = '(%s) ' % prompt
116        create_if_needed(env_dir)
117        executable = sys._base_executable
118        dirname, exename = os.path.split(os.path.abspath(executable))
119        context.executable = executable
120        context.python_dir = dirname
121        context.python_exe = exename
122        if sys.platform == 'win32':
123            binname = 'Scripts'
124            incpath = 'Include'
125            libpath = os.path.join(env_dir, 'Lib', 'site-packages')
126        else:
127            binname = 'bin'
128            incpath = 'include'
129            libpath = os.path.join(env_dir, 'lib',
130                                   'python%d.%d' % sys.version_info[:2],
131                                   'site-packages')
132        context.inc_path = path = os.path.join(env_dir, incpath)
133        create_if_needed(path)
134        create_if_needed(libpath)
135        # Issue 21197: create lib64 as a symlink to lib on 64-bit non-OS X POSIX
136        if ((sys.maxsize > 2**32) and (os.name == 'posix') and
137            (sys.platform != 'darwin')):
138            link_path = os.path.join(env_dir, 'lib64')
139            if not os.path.exists(link_path):   # Issue #21643
140                os.symlink('lib', link_path)
141        context.bin_path = binpath = os.path.join(env_dir, binname)
142        context.bin_name = binname
143        context.env_exe = os.path.join(binpath, exename)
144        create_if_needed(binpath)
145        # Assign and update the command to use when launching the newly created
146        # environment, in case it isn't simply the executable script (e.g. bpo-45337)
147        context.env_exec_cmd = context.env_exe
148        if sys.platform == 'win32':
149            # bpo-45337: Fix up env_exec_cmd to account for file system redirections.
150            # Some redirects only apply to CreateFile and not CreateProcess
151            real_env_exe = os.path.realpath(context.env_exe)
152            if os.path.normcase(real_env_exe) != os.path.normcase(context.env_exe):
153                logger.warning('Actual environment location may have moved due to '
154                               'redirects, links or junctions.\n'
155                               '  Requested location: "%s"\n'
156                               '  Actual location:    "%s"',
157                               context.env_exe, real_env_exe)
158                context.env_exec_cmd = real_env_exe
159        return context
160
161    def create_configuration(self, context):
162        """
163        Create a configuration file indicating where the environment's Python
164        was copied from, and whether the system site-packages should be made
165        available in the environment.
166
167        :param context: The information for the environment creation request
168                        being processed.
169        """
170        context.cfg_path = path = os.path.join(context.env_dir, 'pyvenv.cfg')
171        with open(path, 'w', encoding='utf-8') as f:
172            f.write('home = %s\n' % context.python_dir)
173            if self.system_site_packages:
174                incl = 'true'
175            else:
176                incl = 'false'
177            f.write('include-system-site-packages = %s\n' % incl)
178            f.write('version = %d.%d.%d\n' % sys.version_info[:3])
179            if self.prompt is not None:
180                f.write(f'prompt = {self.prompt!r}\n')
181
182    if os.name != 'nt':
183        def symlink_or_copy(self, src, dst, relative_symlinks_ok=False):
184            """
185            Try symlinking a file, and if that fails, fall back to copying.
186            """
187            force_copy = not self.symlinks
188            if not force_copy:
189                try:
190                    if not os.path.islink(dst): # can't link to itself!
191                        if relative_symlinks_ok:
192                            assert os.path.dirname(src) == os.path.dirname(dst)
193                            os.symlink(os.path.basename(src), dst)
194                        else:
195                            os.symlink(src, dst)
196                except Exception:   # may need to use a more specific exception
197                    logger.warning('Unable to symlink %r to %r', src, dst)
198                    force_copy = True
199            if force_copy:
200                shutil.copyfile(src, dst)
201    else:
202        def symlink_or_copy(self, src, dst, relative_symlinks_ok=False):
203            """
204            Try symlinking a file, and if that fails, fall back to copying.
205            """
206            bad_src = os.path.lexists(src) and not os.path.exists(src)
207            if self.symlinks and not bad_src and not os.path.islink(dst):
208                try:
209                    if relative_symlinks_ok:
210                        assert os.path.dirname(src) == os.path.dirname(dst)
211                        os.symlink(os.path.basename(src), dst)
212                    else:
213                        os.symlink(src, dst)
214                    return
215                except Exception:   # may need to use a more specific exception
216                    logger.warning('Unable to symlink %r to %r', src, dst)
217
218            # On Windows, we rewrite symlinks to our base python.exe into
219            # copies of venvlauncher.exe
220            basename, ext = os.path.splitext(os.path.basename(src))
221            srcfn = os.path.join(os.path.dirname(__file__),
222                                 "scripts",
223                                 "nt",
224                                 basename + ext)
225            # Builds or venv's from builds need to remap source file
226            # locations, as we do not put them into Lib/venv/scripts
227            if sysconfig.is_python_build(True) or not os.path.isfile(srcfn):
228                if basename.endswith('_d'):
229                    ext = '_d' + ext
230                    basename = basename[:-2]
231                if basename == 'python':
232                    basename = 'venvlauncher'
233                elif basename == 'pythonw':
234                    basename = 'venvwlauncher'
235                src = os.path.join(os.path.dirname(src), basename + ext)
236            else:
237                src = srcfn
238            if not os.path.exists(src):
239                if not bad_src:
240                    logger.warning('Unable to copy %r', src)
241                return
242
243            shutil.copyfile(src, dst)
244
245    def setup_python(self, context):
246        """
247        Set up a Python executable in the environment.
248
249        :param context: The information for the environment creation request
250                        being processed.
251        """
252        binpath = context.bin_path
253        path = context.env_exe
254        copier = self.symlink_or_copy
255        dirname = context.python_dir
256        if os.name != 'nt':
257            copier(context.executable, path)
258            if not os.path.islink(path):
259                os.chmod(path, 0o755)
260            for suffix in ('python', 'python3', f'python3.{sys.version_info[1]}'):
261                path = os.path.join(binpath, suffix)
262                if not os.path.exists(path):
263                    # Issue 18807: make copies if
264                    # symlinks are not wanted
265                    copier(context.env_exe, path, relative_symlinks_ok=True)
266                    if not os.path.islink(path):
267                        os.chmod(path, 0o755)
268        else:
269            if self.symlinks:
270                # For symlinking, we need a complete copy of the root directory
271                # If symlinks fail, you'll get unnecessary copies of files, but
272                # we assume that if you've opted into symlinks on Windows then
273                # you know what you're doing.
274                suffixes = [
275                    f for f in os.listdir(dirname) if
276                    os.path.normcase(os.path.splitext(f)[1]) in ('.exe', '.dll')
277                ]
278                if sysconfig.is_python_build(True):
279                    suffixes = [
280                        f for f in suffixes if
281                        os.path.normcase(f).startswith(('python', 'vcruntime'))
282                    ]
283            else:
284                suffixes = ['python.exe', 'python_d.exe', 'pythonw.exe',
285                            'pythonw_d.exe']
286
287            for suffix in suffixes:
288                src = os.path.join(dirname, suffix)
289                if os.path.lexists(src):
290                    copier(src, os.path.join(binpath, suffix))
291
292            if sysconfig.is_python_build(True):
293                # copy init.tcl
294                for root, dirs, files in os.walk(context.python_dir):
295                    if 'init.tcl' in files:
296                        tcldir = os.path.basename(root)
297                        tcldir = os.path.join(context.env_dir, 'Lib', tcldir)
298                        if not os.path.exists(tcldir):
299                            os.makedirs(tcldir)
300                        src = os.path.join(root, 'init.tcl')
301                        dst = os.path.join(tcldir, 'init.tcl')
302                        shutil.copyfile(src, dst)
303                        break
304
305    def _setup_pip(self, context):
306        """Installs or upgrades pip in a virtual environment"""
307        # We run ensurepip in isolated mode to avoid side effects from
308        # environment vars, the current directory and anything else
309        # intended for the global Python environment
310        cmd = [context.env_exec_cmd, '-Im', 'ensurepip', '--upgrade',
311                                                         '--default-pip']
312        subprocess.check_output(cmd, stderr=subprocess.STDOUT)
313
314    def setup_scripts(self, context):
315        """
316        Set up scripts into the created environment from a directory.
317
318        This method installs the default scripts into the environment
319        being created. You can prevent the default installation by overriding
320        this method if you really need to, or if you need to specify
321        a different location for the scripts to install. By default, the
322        'scripts' directory in the venv package is used as the source of
323        scripts to install.
324        """
325        path = os.path.abspath(os.path.dirname(__file__))
326        path = os.path.join(path, 'scripts')
327        self.install_scripts(context, path)
328
329    def post_setup(self, context):
330        """
331        Hook for post-setup modification of the venv. Subclasses may install
332        additional packages or scripts here, add activation shell scripts, etc.
333
334        :param context: The information for the environment creation request
335                        being processed.
336        """
337        pass
338
339    def replace_variables(self, text, context):
340        """
341        Replace variable placeholders in script text with context-specific
342        variables.
343
344        Return the text passed in , but with variables replaced.
345
346        :param text: The text in which to replace placeholder variables.
347        :param context: The information for the environment creation request
348                        being processed.
349        """
350        text = text.replace('__VENV_DIR__', context.env_dir)
351        text = text.replace('__VENV_NAME__', context.env_name)
352        text = text.replace('__VENV_PROMPT__', context.prompt)
353        text = text.replace('__VENV_BIN_NAME__', context.bin_name)
354        text = text.replace('__VENV_PYTHON__', context.env_exe)
355        return text
356
357    def install_scripts(self, context, path):
358        """
359        Install scripts into the created environment from a directory.
360
361        :param context: The information for the environment creation request
362                        being processed.
363        :param path:    Absolute pathname of a directory containing script.
364                        Scripts in the 'common' subdirectory of this directory,
365                        and those in the directory named for the platform
366                        being run on, are installed in the created environment.
367                        Placeholder variables are replaced with environment-
368                        specific values.
369        """
370        binpath = context.bin_path
371        plen = len(path)
372        for root, dirs, files in os.walk(path):
373            if root == path: # at top-level, remove irrelevant dirs
374                for d in dirs[:]:
375                    if d not in ('common', os.name):
376                        dirs.remove(d)
377                continue # ignore files in top level
378            for f in files:
379                if (os.name == 'nt' and f.startswith('python')
380                        and f.endswith(('.exe', '.pdb'))):
381                    continue
382                srcfile = os.path.join(root, f)
383                suffix = root[plen:].split(os.sep)[2:]
384                if not suffix:
385                    dstdir = binpath
386                else:
387                    dstdir = os.path.join(binpath, *suffix)
388                if not os.path.exists(dstdir):
389                    os.makedirs(dstdir)
390                dstfile = os.path.join(dstdir, f)
391                with open(srcfile, 'rb') as f:
392                    data = f.read()
393                if not srcfile.endswith(('.exe', '.pdb')):
394                    try:
395                        data = data.decode('utf-8')
396                        data = self.replace_variables(data, context)
397                        data = data.encode('utf-8')
398                    except UnicodeError as e:
399                        data = None
400                        logger.warning('unable to copy script %r, '
401                                       'may be binary: %s', srcfile, e)
402                if data is not None:
403                    with open(dstfile, 'wb') as f:
404                        f.write(data)
405                    shutil.copymode(srcfile, dstfile)
406
407    def upgrade_dependencies(self, context):
408        logger.debug(
409            f'Upgrading {CORE_VENV_DEPS} packages in {context.bin_path}'
410        )
411        cmd = [context.env_exec_cmd, '-m', 'pip', 'install', '--upgrade']
412        cmd.extend(CORE_VENV_DEPS)
413        subprocess.check_call(cmd)
414
415
416def create(env_dir, system_site_packages=False, clear=False,
417           symlinks=False, with_pip=False, prompt=None, upgrade_deps=False):
418    """Create a virtual environment in a directory."""
419    builder = EnvBuilder(system_site_packages=system_site_packages,
420                         clear=clear, symlinks=symlinks, with_pip=with_pip,
421                         prompt=prompt, upgrade_deps=upgrade_deps)
422    builder.create(env_dir)
423
424def main(args=None):
425    compatible = True
426    if sys.version_info < (3, 3):
427        compatible = False
428    elif not hasattr(sys, 'base_prefix'):
429        compatible = False
430    if not compatible:
431        raise ValueError('This script is only for use with Python >= 3.3')
432    else:
433        import argparse
434
435        parser = argparse.ArgumentParser(prog=__name__,
436                                         description='Creates virtual Python '
437                                                     'environments in one or '
438                                                     'more target '
439                                                     'directories.',
440                                         epilog='Once an environment has been '
441                                                'created, you may wish to '
442                                                'activate it, e.g. by '
443                                                'sourcing an activate script '
444                                                'in its bin directory.')
445        parser.add_argument('dirs', metavar='ENV_DIR', nargs='+',
446                            help='A directory to create the environment in.')
447        parser.add_argument('--system-site-packages', default=False,
448                            action='store_true', dest='system_site',
449                            help='Give the virtual environment access to the '
450                                 'system site-packages dir.')
451        if os.name == 'nt':
452            use_symlinks = False
453        else:
454            use_symlinks = True
455        group = parser.add_mutually_exclusive_group()
456        group.add_argument('--symlinks', default=use_symlinks,
457                           action='store_true', dest='symlinks',
458                           help='Try to use symlinks rather than copies, '
459                                'when symlinks are not the default for '
460                                'the platform.')
461        group.add_argument('--copies', default=not use_symlinks,
462                           action='store_false', dest='symlinks',
463                           help='Try to use copies rather than symlinks, '
464                                'even when symlinks are the default for '
465                                'the platform.')
466        parser.add_argument('--clear', default=False, action='store_true',
467                            dest='clear', help='Delete the contents of the '
468                                               'environment directory if it '
469                                               'already exists, before '
470                                               'environment creation.')
471        parser.add_argument('--upgrade', default=False, action='store_true',
472                            dest='upgrade', help='Upgrade the environment '
473                                               'directory to use this version '
474                                               'of Python, assuming Python '
475                                               'has been upgraded in-place.')
476        parser.add_argument('--without-pip', dest='with_pip',
477                            default=True, action='store_false',
478                            help='Skips installing or upgrading pip in the '
479                                 'virtual environment (pip is bootstrapped '
480                                 'by default)')
481        parser.add_argument('--prompt',
482                            help='Provides an alternative prompt prefix for '
483                                 'this environment.')
484        parser.add_argument('--upgrade-deps', default=False, action='store_true',
485                            dest='upgrade_deps',
486                            help='Upgrade core dependencies: {} to the latest '
487                                 'version in PyPI'.format(
488                                 ' '.join(CORE_VENV_DEPS)))
489        options = parser.parse_args(args)
490        if options.upgrade and options.clear:
491            raise ValueError('you cannot supply --upgrade and --clear together.')
492        builder = EnvBuilder(system_site_packages=options.system_site,
493                             clear=options.clear,
494                             symlinks=options.symlinks,
495                             upgrade=options.upgrade,
496                             with_pip=options.with_pip,
497                             prompt=options.prompt,
498                             upgrade_deps=options.upgrade_deps)
499        for d in options.dirs:
500            builder.create(d)
501
502if __name__ == '__main__':
503    rc = 1
504    try:
505        main()
506        rc = 0
507    except Exception as e:
508        print('Error: %s' % e, file=sys.stderr)
509    sys.exit(rc)
510