Globally manage and track some temp build directories

src/pip/_internal/build_env.py

@@ -17,7 +17,7 @@
 
 from pip import __file__ as pip_location
 from pip._internal.utils.subprocess import call_subprocess
-from pip._internal.utils.temp_dir import TempDirectory
+from pip._internal.utils.temp_dir import TempDirectory, tempdir_kinds
 from pip._internal.utils.typing import MYPY_CHECK_RUNNING
 from pip._internal.utils.ui import open_spinner
 
@@ -54,10 +54,12 @@ class BuildEnvironment(object):
 
     def __init__(self):
         # type: () -> None
-        self._temp_dir = TempDirectory(kind="build-env")
+        temp_dir = TempDirectory(
+            kind=tempdir_kinds.BUILD_ENV, globally_managed=True
+        )
 
         self._prefixes = OrderedDict((
-            (name, _Prefix(os.path.join(self._temp_dir.path, name)))
+            (name, _Prefix(os.path.join(temp_dir.path, name)))
             for name in ('normal', 'overlay')
         ))
 
@@ -76,7 +78,7 @@ def __init__(self):
                 get_python_lib(plat_specific=True),
             )
         }
-        self._site_dir = os.path.join(self._temp_dir.path, 'site')
+        self._site_dir = os.path.join(temp_dir.path, 'site')
         if not os.path.exists(self._site_dir):
             os.mkdir(self._site_dir)
         with open(os.path.join(self._site_dir, 'sitecustomize.py'), 'w') as fp:
@@ -133,10 +135,6 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             else:
                 os.environ[varname] = old_value
 
-    def cleanup(self):
-        # type: () -> None
-        self._temp_dir.cleanup()
-
     def check_requirements(self, reqs):
         # type: (Iterable[str]) -> Tuple[Set[Tuple[str, str]], Set[str]]
         """Return 2 sets:

src/pip/_internal/cli/req_command.py

@@ -9,9 +9,10 @@
 import os
 from functools import partial
 
+from pip._internal.cli import cmdoptions
 from pip._internal.cli.base_command import Command
 from pip._internal.cli.command_context import CommandContextMixIn
-from pip._internal.exceptions import CommandError
+from pip._internal.exceptions import CommandError, PreviousBuildDirError
 from pip._internal.index.package_finder import PackageFinder
 from pip._internal.legacy_resolve import Resolver
 from pip._internal.models.selection_prefs import SelectionPreferences
@@ -28,16 +29,22 @@
     make_link_collector,
     pip_self_version_check,
 )
+from pip._internal.utils.temp_dir import tempdir_kinds
 from pip._internal.utils.typing import MYPY_CHECK_RUNNING
 
 if MYPY_CHECK_RUNNING:
     from optparse import Values
-    from typing import List, Optional, Tuple
+    from typing import Any, List, Optional, Tuple
+
     from pip._internal.cache import WheelCache
     from pip._internal.models.target_python import TargetPython
     from pip._internal.req.req_set import RequirementSet
     from pip._internal.req.req_tracker import RequirementTracker
-    from pip._internal.utils.temp_dir import TempDirectory
+    from pip._internal.utils.temp_dir import (
+        TempDirectory,
+        TempDirectoryTypeRegistry,
+    )
+
 
 logger = logging.getLogger(__name__)
 
@@ -149,8 +156,45 @@ def handle_pip_version_check(self, options):
             pip_self_version_check(session, options)
 
 
+KEEPABLE_TEMPDIR_TYPES = [tempdir_kinds.BUILD_ENV, tempdir_kinds.REQ_BUILD]
+
+
+def with_cleanup(func):
+    # type: (Any) -> Any
+    """Decorator for common logic related to managing temporary
+    directories.
+    """
+    def configure_tempdir_registry(registry):
+        # type: (TempDirectoryTypeRegistry) -> None
+        for t in KEEPABLE_TEMPDIR_TYPES:
+            registry.set_delete(t, False)
+
+    def wrapper(self, options, args):
+        # type: (RequirementCommand, Values, List[Any]) -> Optional[int]
+        assert self.tempdir_registry is not None
+        if options.no_clean:
+            configure_tempdir_registry(self.tempdir_registry)
+
+        try:
+            return func(self, options, args)
+        except PreviousBuildDirError:
+            # This kind of conflict can occur when the user passes an explicit
+            # build directory with a pre-existing folder. In that case we do
+            # not want to accidentally remove it.
+            configure_tempdir_registry(self.tempdir_registry)
+            raise
+
+    return wrapper
+
+
 class RequirementCommand(IndexGroupCommand):
 
+    def __init__(self, *args, **kw):
+        # type: (Any, Any) -> None
+        super(RequirementCommand, self).__init__(*args, **kw)
+
+        self.cmd_opts.add_option(cmdoptions.no_clean())
+
     @staticmethod
     def make_requirement_preparer(
         temp_build_dir,           # type: TempDirectory

src/pip/_internal/commands/download.py

@@ -8,7 +8,7 @@
 
 from pip._internal.cli import cmdoptions
 from pip._internal.cli.cmdoptions import make_target_python
-from pip._internal.cli.req_command import RequirementCommand
+from pip._internal.cli.req_command import RequirementCommand, with_cleanup
 from pip._internal.req import RequirementSet
 from pip._internal.req.req_tracker import get_requirement_tracker
 from pip._internal.utils.misc import ensure_dir, normalize_path, write_output
@@ -52,7 +52,6 @@ def __init__(self, *args, **kw):
         cmd_opts.add_option(cmdoptions.prefer_binary())
         cmd_opts.add_option(cmdoptions.src())
         cmd_opts.add_option(cmdoptions.pre())
-        cmd_opts.add_option(cmdoptions.no_clean())
         cmd_opts.add_option(cmdoptions.require_hashes())
         cmd_opts.add_option(cmdoptions.progress_bar())
         cmd_opts.add_option(cmdoptions.no_build_isolation())
@@ -77,6 +76,7 @@ def __init__(self, *args, **kw):
         self.parser.insert_option_group(0, index_opts)
         self.parser.insert_option_group(0, cmd_opts)
 
+    @with_cleanup
     def run(self, options, args):
         options.ignore_installed = True
         # editable doesn't really make sense for `pip download`, but the bowels

src/pip/_internal/commands/install.py

@@ -21,7 +21,7 @@
 from pip._internal.cache import WheelCache
 from pip._internal.cli import cmdoptions
 from pip._internal.cli.cmdoptions import make_target_python
-from pip._internal.cli.req_command import RequirementCommand
+from pip._internal.cli.req_command import RequirementCommand, with_cleanup
 from pip._internal.cli.status_codes import ERROR, SUCCESS
 from pip._internal.exceptions import (
     CommandError,
@@ -227,7 +227,6 @@ def __init__(self, *args, **kw):
         cmd_opts.add_option(cmdoptions.no_binary())
         cmd_opts.add_option(cmdoptions.only_binary())
         cmd_opts.add_option(cmdoptions.prefer_binary())
-        cmd_opts.add_option(cmdoptions.no_clean())
         cmd_opts.add_option(cmdoptions.require_hashes())
         cmd_opts.add_option(cmdoptions.progress_bar())
 
@@ -239,6 +238,7 @@ def __init__(self, *args, **kw):
         self.parser.insert_option_group(0, index_opts)
         self.parser.insert_option_group(0, cmd_opts)
 
+    @with_cleanup
     def run(self, options, args):
         # type: (Values, List[Any]) -> int
         cmdoptions.check_install_build_global(options)

src/pip/_internal/commands/wheel.py

@@ -11,7 +11,7 @@
 
 from pip._internal.cache import WheelCache
 from pip._internal.cli import cmdoptions
-from pip._internal.cli.req_command import RequirementCommand
+from pip._internal.cli.req_command import RequirementCommand, with_cleanup
 from pip._internal.exceptions import CommandError, PreviousBuildDirError
 from pip._internal.req import RequirementSet
 from pip._internal.req.req_tracker import get_requirement_tracker
@@ -101,7 +101,6 @@ def __init__(self, *args, **kw):
                   "pip only finds stable versions."),
         )
 
-        cmd_opts.add_option(cmdoptions.no_clean())
         cmd_opts.add_option(cmdoptions.require_hashes())
 
         index_opts = cmdoptions.make_option_group(
@@ -112,6 +111,7 @@ def __init__(self, *args, **kw):
         self.parser.insert_option_group(0, index_opts)
         self.parser.insert_option_group(0, cmd_opts)
 
+    @with_cleanup
     def run(self, options, args):
         # type: (Values, List[Any]) -> None
         cmdoptions.check_install_build_global(options)

src/pip/_internal/req/req_install.py

@@ -49,7 +49,7 @@
     rmtree,
 )
 from pip._internal.utils.packaging import get_metadata
-from pip._internal.utils.temp_dir import TempDirectory
+from pip._internal.utils.temp_dir import TempDirectory, tempdir_kinds
 from pip._internal.utils.typing import MYPY_CHECK_RUNNING
 from pip._internal.utils.virtualenv import running_under_virtualenv
 from pip._internal.vcs import vcs
@@ -358,7 +358,9 @@ def ensure_build_location(self, build_dir):
             # Some systems have /tmp as a symlink which confuses custom
             # builds (such as numpy). Thus, we ensure that the real path
             # is returned.
-            self._temp_build_dir = TempDirectory(kind="req-build")
+            self._temp_build_dir = TempDirectory(
+                kind=tempdir_kinds.REQ_BUILD, globally_managed=True
+            )
 
             return self._temp_build_dir.path
         if self.editable:
@@ -418,10 +420,7 @@ def remove_temporary_source(self):
             logger.debug('Removing source in %s', self.source_dir)
             rmtree(self.source_dir)
         self.source_dir = None
-        if self._temp_build_dir:
-            self._temp_build_dir.cleanup()
-            self._temp_build_dir = None
-        self.build_env.cleanup()
+        self._temp_build_dir = None
 
     def check_if_exists(self, use_user_site):
         # type: (bool) -> None

src/pip/_internal/utils/temp_dir.py

@@ -9,7 +9,7 @@
 
 from pip._vendor.contextlib2 import ExitStack
 
-from pip._internal.utils.misc import rmtree
+from pip._internal.utils.misc import enum, rmtree
 from pip._internal.utils.typing import MYPY_CHECK_RUNNING
 
 if MYPY_CHECK_RUNNING:
@@ -21,6 +21,14 @@
 logger = logging.getLogger(__name__)
 
 
+# Kinds of temporary directories. Only needed for ones that are
+# globally-managed.
+tempdir_kinds = enum(
+    BUILD_ENV="build-env",
+    REQ_BUILD="req-build",
+)
+
+
 _tempdir_manager = None  # type: Optional[ExitStack]
 
 

tests/unit/test_build_env.py

@@ -29,6 +29,7 @@ def run_with_build_env(script, setup_script_contents,
                 SelectionPreferences
             )
             from pip._internal.network.session import PipSession
+            from pip._internal.utils.temp_dir import global_tempdir_manager
 
             link_collector = LinkCollector(
                 session=PipSession(),
@@ -41,19 +42,21 @@ def run_with_build_env(script, setup_script_contents,
                 link_collector=link_collector,
                 selection_prefs=selection_prefs,
             )
-            build_env = BuildEnvironment()
 
-            try:
+            with global_tempdir_manager():
+                build_env = BuildEnvironment()
             ''' % str(script.scratch_path)) +
         indent(dedent(setup_script_contents), '    ') +
-        dedent(
-            '''
+        indent(
+            dedent(
+                '''
                 if len(sys.argv) > 1:
                     with build_env:
                         subprocess.check_call((sys.executable, sys.argv[1]))
-            finally:
-                build_env.cleanup()
-            ''')
+                '''
+            ),
+            '    '
+        )
     )
     args = ['python', build_env_script]
     if test_script_contents is not None: