Use tempdir_registry to control auto-deleted files globally

Next we can actually transition some files to be globally-managed that
are only deleted conditionally.

Author: chrahunt

src/pip/_internal/cli/base_command.py

@@ -34,14 +34,21 @@
 from pip._internal.utils.filesystem import check_path_owner
 from pip._internal.utils.logging import BrokenStdoutLoggingError, setup_logging
 from pip._internal.utils.misc import get_prog, normalize_path
-from pip._internal.utils.temp_dir import global_tempdir_manager
+from pip._internal.utils.temp_dir import (
+    global_tempdir_manager,
+    tempdir_registry,
+)
 from pip._internal.utils.typing import MYPY_CHECK_RUNNING
 from pip._internal.utils.virtualenv import running_under_virtualenv
 
 if MYPY_CHECK_RUNNING:
-    from typing import List, Tuple, Any
+    from typing import List, Optional, Tuple, Any
     from optparse import Values
 
+    from pip._internal.utils.temp_dir import (
+        TempDirectoryTypeRegistry as TempDirRegistry
+    )
+
 __all__ = ['Command']
 
 logger = logging.getLogger(__name__)
@@ -68,6 +75,8 @@ def __init__(self, name, summary, isolated=False):
         self.summary = summary
         self.parser = ConfigOptionParser(**parser_kw)
 
+        self.tempdir_registry = None  # type: Optional[TempDirRegistry]
+
         # Commands should add options to this option group
         optgroup_name = '%s Options' % self.name.capitalize()
         self.cmd_opts = optparse.OptionGroup(self.parser, optgroup_name)
@@ -108,6 +117,10 @@ def main(self, args):
 
     def _main(self, args):
         # type: (List[str]) -> int
+        # We must initialize this before the tempdir manager, otherwise the
+        # configuration would not be accessible by the time we clean up the
+        # tempdir manager.
+        self.tempdir_registry = self.enter_context(tempdir_registry())
         # Intentionally set as early as possible so globally-managed temporary
         # directories are available to the rest of the code.
         self.enter_context(global_tempdir_manager())

tests/conftest.py

@@ -10,7 +10,7 @@
 
 import pytest
 import six
-from pip._vendor.contextlib2 import ExitStack
+from pip._vendor.contextlib2 import ExitStack, nullcontext
 from setuptools.wheel import Wheel
 
 from pip._internal.cli.main import main as pip_entry_point
@@ -188,13 +188,18 @@ def isolate(tmpdir):
 
 
 @pytest.fixture(autouse=True)
-def scoped_global_tempdir_manager():
+def scoped_global_tempdir_manager(request):
     """Make unit tests with globally-managed tempdirs easier
 
     Each test function gets its own individual scope for globally-managed
     temporary directories in the application.
     """
-    with global_tempdir_manager():
+    if "no_auto_tempdir_manager" in request.keywords:
+        ctx = nullcontext
+    else:
+        ctx = global_tempdir_manager
+
+    with ctx():
         yield
 
 

tests/unit/test_base_command.py

@@ -2,10 +2,14 @@
 import os
 import time
 
-from mock import patch
+import pytest
+from mock import Mock, patch
 
 from pip._internal.cli.base_command import Command
+from pip._internal.cli.status_codes import SUCCESS
+from pip._internal.utils import temp_dir
 from pip._internal.utils.logging import BrokenStdoutLoggingError
+from pip._internal.utils.temp_dir import TempDirectory
 
 
 class FakeCommand(Command):
@@ -145,3 +149,65 @@ def test_unicode_messages(self, tmpdir):
         cmd = FakeCommandWithUnicode()
         log_path = tmpdir.joinpath('log')
         cmd.main(['fake_unicode', '--log', log_path])
+
+
+@pytest.mark.no_auto_tempdir_manager
+def test_base_command_provides_tempdir_helpers():
+    assert temp_dir._tempdir_manager is None
+    assert temp_dir._tempdir_registry is None
+
+    def assert_helpers_set(options, args):
+        assert temp_dir._tempdir_manager is not None
+        assert temp_dir._tempdir_registry is not None
+
+    c = Command("fake", "fake")
+    c.run = Mock(side_effect=assert_helpers_set)
+    assert c.main(["fake"]) == SUCCESS
+    c.run.assert_called_once()
+
+
+not_deleted = "not_deleted"
+
+
+@pytest.mark.parametrize("kind,exists", [
+    (not_deleted, True), ("deleted", False)
+])
+@pytest.mark.no_auto_tempdir_manager
+def test_base_command_global_tempdir_cleanup(kind, exists):
+    assert temp_dir._tempdir_manager is None
+    assert temp_dir._tempdir_registry is None
+
+    class Holder(object):
+        value = None
+
+    def create_temp_dirs(options, args):
+        c.tempdir_registry.set_delete(not_deleted, False)
+        Holder.value = TempDirectory(kind=kind, globally_managed=True).path
+
+    c = Command("fake", "fake")
+    c.run = Mock(side_effect=create_temp_dirs)
+    assert c.main(["fake"]) == SUCCESS
+    c.run.assert_called_once()
+    assert os.path.exists(Holder.value) == exists
+
+
+@pytest.mark.parametrize("kind,exists", [
+    (not_deleted, True), ("deleted", False)
+])
+@pytest.mark.no_auto_tempdir_manager
+def test_base_command_local_tempdir_cleanup(kind, exists):
+    assert temp_dir._tempdir_manager is None
+    assert temp_dir._tempdir_registry is None
+
+    def create_temp_dirs(options, args):
+        c.tempdir_registry.set_delete(not_deleted, False)
+
+        with TempDirectory(kind=kind) as d:
+            path = d.path
+            assert os.path.exists(path)
+        assert os.path.exists(path) == exists
+
+    c = Command("fake", "fake")
+    c.run = Mock(side_effect=create_temp_dirs)
+    assert c.main(["fake"]) == SUCCESS
+    c.run.assert_called_once()