Add --upgrade-deps to venv module

- This allows for pip + setuptools to be automatically upgraded to the latest version on PyPI - Update documentation to represent this change bpo-34556: Add --upgrade to venv module
ambv · Jun 17, 2019 · Jun 17, 2019 · May 5, 2019 · May 5, 2019 · May 5, 2019
commit 42b6f1abbdb560759e22eb25df3c3e653a8827a7
diff --git a/Doc/library/venv.rst b/Doc/library/venv.rst
@@ -97,7 +97,7 @@ creation according to their needs, the :class:`EnvBuilder` class.
 
 .. class:: EnvBuilder(system_site_packages=False, clear=False, \
                       symlinks=False, upgrade=False, with_pip=False, \
-                      prompt=None)
+                      prompt=None, upgrade_deps=False)
 
     The :class:`EnvBuilder` class accepts the following keyword arguments on
     instantiation:
@@ -123,12 +123,17 @@ creation according to their needs, the :class:`EnvBuilder` class.
       (defaults to ``None`` which means directory name of the environment would
       be used).
 
+    * ``upgrade_deps`` - Update the base venv modules to the latest on PyPI
+
     .. versionchanged:: 3.4
        Added the ``with_pip`` parameter
 
     .. versionadded:: 3.6
        Added the ``prompt`` parameter
 
+    .. versionadded:: 3.8
+       Added the ``upgrade_deps`` parameter
+
     Creators of third-party virtual environment tools will be free to use the
     provided ``EnvBuilder`` class as a base class.
 

diff --git a/Doc/using/venv-create.inc b/Doc/using/venv-create.inc
@@ -35,7 +35,7 @@ your :ref:`Python installation <using-on-windows>`::
 The command, if run with ``-h``, will show the available options::
 
     usage: venv [-h] [--system-site-packages] [--symlinks | --copies] [--clear]
-                [--upgrade] [--without-pip] [--prompt PROMPT]
+                [--upgrade] [--without-pip] [--prompt PROMPT] [--upgrade-deps]
                 ENV_DIR [ENV_DIR ...]
 
     Creates virtual Python environments in one or more target directories.
@@ -60,10 +60,18 @@ The command, if run with ``-h``, will show the available options::
                             environment (pip is bootstrapped by default)
       --prompt PROMPT       Provides an alternative prompt prefix for this
                             environment.
+      --upgrade-deps        Upgrade core dependencies: pip setuptools to the
+                            latest version in PyPI
 
     Once an environment has been created, you may wish to activate it, e.g. by
     sourcing an activate script in its bin directory.
 
+.. versionchanged:: 3.8
+   Add ``--upgrade-deps`` option to upgrade pip + setuptools to the latest on PyPI
+
+.. versionchanged:: 3.6
+   Add ability to change the activated venv's prompt via ``--prompt``
+
 .. versionchanged:: 3.4
    Installs pip by default, added the ``--without-pip``  and ``--copies``
    options

diff --git a/Lib/test/test_venv.py b/Lib/test/test_venv.py
@@ -16,9 +16,9 @@
 from test.support import (captured_stdout, captured_stderr, requires_zlib,
                           can_symlink, EnvironmentVarGuard, rmtree,
                           import_module)
-import threading
 import unittest
 import venv
+from unittest.mock import patch
 
 try:
     import ctypes
@@ -32,17 +32,21 @@
     or sys.prefix == sys.base_prefix,
     'cannot run venv.create from within a venv on this platform')
 
+
 def check_output(cmd, encoding=None):
-    p = subprocess.Popen(cmd,
+    p = subprocess.Popen(
+        cmd,
         stdout=subprocess.PIPE,
         stderr=subprocess.PIPE,
-        encoding=encoding)
+        encoding=encoding
+    )
     out, err = p.communicate()
     if p.returncode:
         raise subprocess.CalledProcessError(
             p.returncode, cmd, out, err)
     return out, err
 
+
 class BaseTest(unittest.TestCase):
     """Base class for venv tests."""
     maxDiff = 80 * 50
@@ -77,6 +81,7 @@ def get_text_file_contents(self, *args):
             result = f.read()
         return result
 
+
 class BasicTest(BaseTest):
     """Test venv module functionality."""
 
@@ -131,12 +136,33 @@ def test_prompt(self):
         self.assertEqual(context.prompt, '(My prompt) ')
         self.assertIn("prompt = 'My prompt'\n", data)
 
+    def test_upgrade_dependencies(self):
+        builder = venv.EnvBuilder()
+        pip_exe = 'pip.exe' if sys.platform == 'win32' else 'pip'
+        with tempfile.TemporaryDirectory() as fake_env_dir:
+
+            def pip_cmd_checker(cmd):
+                self.assertEqual(
+                    cmd,
+                    [
+                        os.path.join(fake_env_dir, pip_exe),
+                        'install',
+                        '-U',
+                        'pip',
+                        'setuptools'
+                    ]
+                )
+
+            fake_context = builder.ensure_directories(fake_env_dir)
+            with patch("venv.subprocess.check_call", pip_cmd_checker):
+                builder.upgrade_dependencies(fake_context, fake_env_dir)
+
     @requireVenvCreate
     def test_prefixes(self):
         """
         Test that the prefix values are as expected.
         """
-        #check our prefixes
+        # check our prefixes
         self.assertEqual(sys.base_prefix, sys.prefix)
         self.assertEqual(sys.base_exec_prefix, sys.exec_prefix)
 
@@ -149,7 +175,8 @@ def test_prefixes(self):
             ('prefix', self.env_dir),
             ('prefix', self.env_dir),
             ('base_prefix', sys.prefix),
-            ('base_exec_prefix', sys.exec_prefix)):
+            ('base_exec_prefix', sys.exec_prefix)
+        ):
             cmd[2] = 'import sys; print(sys.%s)' % prefix
             out, err = check_output(cmd)
             self.assertEqual(out.strip(), expected.encode())
@@ -209,7 +236,7 @@ def clear_directory(self, path):
                 rmtree(fn)
 
     def test_unoverwritable_fails(self):
-        #create a file clashing with directories in the env dir
+        # create a file clashing with directories in the env dir
         for paths in self.ENV_SUBDIRS[:3]:
             fn = os.path.join(self.env_dir, *paths)
             with open(fn, 'wb') as f:
@@ -304,7 +331,6 @@ def test_unicode_in_batch_file(self):
         builder = venv.EnvBuilder(clear=True)
         builder.create(env_dir)
         activate = os.path.join(env_dir, self.bindir, 'activate.bat')
-        envpy = os.path.join(env_dir, self.bindir, self.exe)
         out, err = check_output(
             [activate, '&', self.exe, '-c', 'print(0)'],
             encoding='oem',
@@ -466,5 +492,6 @@ def test_with_pip(self):
         self.do_test_with_pip(False)
         self.do_test_with_pip(True)
 
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/Lib/venv/__init__.py b/Lib/venv/__init__.py
@@ -12,6 +12,8 @@
 import sysconfig
 import types
 
+
+CORE_VENV_DEPS = ('pip', 'setuptools')
 logger = logging.getLogger(__name__)
 
 
@@ -38,16 +40,19 @@ class EnvBuilder:
     :param with_pip: If True, ensure pip is installed in the virtual
                      environment
     :param prompt: Alternative terminal prefix for the environment.
+    :param upgrade_deps: Update the base venv modules to the latest on PyPI
     """
 
     def __init__(self, system_site_packages=False, clear=False,
-                 symlinks=False, upgrade=False, with_pip=False, prompt=None):
+                 symlinks=False, upgrade=False, with_pip=False, prompt=None,
+                 upgrade_deps=False):
         self.system_site_packages = system_site_packages
         self.clear = clear
         self.symlinks = symlinks
         self.upgrade = upgrade
         self.with_pip = with_pip
         self.prompt = prompt
+        self.upgrade_deps = upgrade_deps
 
     def create(self, env_dir):
         """
@@ -74,6 +79,8 @@ def create(self, env_dir):
             # restore it and rewrite the configuration
             self.system_site_packages = True
             self.create_configuration(context)
+        if self.upgrade_deps:
+            self.upgrade_dependencies(context, env_dir)
 
     def clear_directory(self, path):
         for fn in os.listdir(path):
@@ -105,7 +112,6 @@ def create_if_needed(d):
         prompt = self.prompt if self.prompt is not None else context.env_name
         context.prompt = '(%s) ' % prompt
         create_if_needed(env_dir)
-        env = os.environ
         executable = getattr(sys, '_base_executable', sys.executable)
         dirname, exename = os.path.split(os.path.abspath(executable))
         context.executable = executable
@@ -250,7 +256,7 @@ def setup_python(self, context):
 
             if sysconfig.is_python_build(True):
                 # copy init.tcl
-                for root, dirs, files in os.walk(context.python_dir):
+                for root, _dirs, files in os.walk(context.python_dir):
                     if 'init.tcl' in files:
                         tcldir = os.path.basename(root)
                         tcldir = os.path.join(context.env_dir, 'Lib', tcldir)
@@ -329,11 +335,11 @@ def install_scripts(self, context, path):
         binpath = context.bin_path
         plen = len(path)
         for root, dirs, files in os.walk(path):
-            if root == path: # at top-level, remove irrelevant dirs
+            if root == path:  # at top-level, remove irrelevant dirs
                 for d in dirs[:]:
                     if d not in ('common', os.name):
                         dirs.remove(d)
-                continue # ignore files in top level
+                continue  # ignore files in top level
             for f in files:
                 if (os.name == 'nt' and f.startswith('python')
                         and f.endswith(('.exe', '.pdb'))):
@@ -363,15 +369,28 @@ def install_scripts(self, context, path):
                         f.write(data)
                     shutil.copymode(srcfile, dstfile)
 
+    def upgrade_dependencies(self, context, env_dir):
+        logging.debug(
+            'Upgrading {} packages in {}'.format(CORE_VENV_DEPS, env_dir)
+        )
+        if sys.platform == 'win32':
+            pip_exe = os.path.join(context.bin_dir, 'pip.exe')
+        else:
+            pip_exe = os.path.join(context.bin_dir, 'pip')
+        cmd = [pip_exe, 'install', '-U']
+        cmd.extend(CORE_VENV_DEPS)
+        subprocess.check_call(cmd)
+
 
 def create(env_dir, system_site_packages=False, clear=False,
-                    symlinks=False, with_pip=False, prompt=None):
+           symlinks=False, with_pip=False, prompt=None):
     """Create a virtual environment in a directory."""
     builder = EnvBuilder(system_site_packages=system_site_packages,
                          clear=clear, symlinks=symlinks, with_pip=with_pip,
                          prompt=prompt)
     builder.create(env_dir)
 
+
 def main(args=None):
     compatible = True
     if sys.version_info < (3, 3):
@@ -432,6 +451,11 @@ def main(args=None):
         parser.add_argument('--prompt',
                             help='Provides an alternative prompt prefix for '
                                  'this environment.')
+        parser.add_argument('--upgrade-deps', default=False, action='store_true',
+                            dest='upgrade_deps',
+                            help='Upgrade core dependencies: {} to the latest '
+                                 'version in PyPI'.format(
+                                 ' '.join(CORE_VENV_DEPS)))
         options = parser.parse_args(args)
         if options.upgrade and options.clear:
             raise ValueError('you cannot supply --upgrade and --clear together.')
@@ -440,10 +464,12 @@ def main(args=None):
                              symlinks=options.symlinks,
                              upgrade=options.upgrade,
                              with_pip=options.with_pip,
-                             prompt=options.prompt)
+                             prompt=options.prompt,
+                             upgrade_deps=options.upgrade_deps)
         for d in options.dirs:
             builder.create(d)
 
+
 if __name__ == '__main__':
     rc = 1
     try: