From 0a57e4e9f27e9e7b2c95bf1c5e6e71a36ecbeded Mon Sep 17 00:00:00 2001
From: Chris Hunt <chrahunt@gmail.com>
Date: Fri, 13 Dec 2019 22:05:22 -0500
Subject: [PATCH] Manage temp directory deletion centrally

This gives us a global toggle that we can use to control whether
temporary directories get deleted from one place (ideally, in the
commands taking --no-clean).
---
 src/pip/_internal/utils/temp_dir.py | 48 +++++++++++++++++++++++++++--
 tests/unit/test_utils_temp_dir.py   | 27 ++++++++++++++++
 2 files changed, 73 insertions(+), 2 deletions(-)

diff --git a/src/pip/_internal/utils/temp_dir.py b/src/pip/_internal/utils/temp_dir.py
index 0be0ff785..2d3422365 100644
--- a/src/pip/_internal/utils/temp_dir.py
+++ b/src/pip/_internal/utils/temp_dir.py
@@ -13,7 +13,7 @@ from pip._internal.utils.misc import rmtree
 from pip._internal.utils.typing import MYPY_CHECK_RUNNING
 
 if MYPY_CHECK_RUNNING:
-    from typing import Any, Iterator, Optional, TypeVar
+    from typing import Any, Dict, Iterator, Optional, TypeVar
 
     _T = TypeVar('_T', bound='TempDirectory')
 
@@ -36,6 +36,47 @@ def global_tempdir_manager():
             _tempdir_manager = old_tempdir_manager
 
 
+class TempDirectoryTypeRegistry(object):
+    """Manages temp directory behavior
+    """
+
+    def __init__(self):
+        # type: () -> None
+        self._should_delete = {}  # type: Dict[str, bool]
+
+    def set_delete(self, kind, value):
+        # type: (str, bool) -> None
+        """Indicate whether a TempDirectory of the given kind should be
+        auto-deleted.
+        """
+        self._should_delete[kind] = value
+
+    def get_delete(self, kind):
+        # type: (str) -> bool
+        """Get configured auto-delete flag for a given TempDirectory type,
+        default True.
+        """
+        return self._should_delete.get(kind, True)
+
+
+_tempdir_registry = None  # type: Optional[TempDirectoryTypeRegistry]
+
+
+@contextmanager
+def tempdir_registry():
+    # type: () -> Iterator[TempDirectoryTypeRegistry]
+    """Provides a scoped global tempdir registry that can be used to dictate
+    whether directories should be deleted.
+    """
+    global _tempdir_registry
+    old_tempdir_registry = _tempdir_registry
+    _tempdir_registry = TempDirectoryTypeRegistry()
+    try:
+        yield _tempdir_registry
+    finally:
+        _tempdir_registry = old_tempdir_registry
+
+
 class TempDirectory(object):
     """Helper class that owns and cleans up a temporary directory.
 
@@ -68,8 +109,11 @@ class TempDirectory(object):
 
         if path is None and delete is None:
             # If we were not given an explicit directory, and we were not given
-            # an explicit delete option, then we'll default to deleting.
+            # an explicit delete option, then we'll default to deleting unless
+            # the tempdir_registry says otherwise.
             delete = True
+            if _tempdir_registry:
+                delete = _tempdir_registry.get_delete(kind)
 
         if path is None:
             path = self._create(kind)
diff --git a/tests/unit/test_utils_temp_dir.py b/tests/unit/test_utils_temp_dir.py
index 9b45a75b9..c395a7861 100644
--- a/tests/unit/test_utils_temp_dir.py
+++ b/tests/unit/test_utils_temp_dir.py
@@ -11,6 +11,7 @@ from pip._internal.utils.temp_dir import (
     AdjacentTempDirectory,
     TempDirectory,
     global_tempdir_manager,
+    tempdir_registry,
 )
 
 
@@ -207,3 +208,29 @@ def test_tempdirectory_asserts_global_tempdir(monkeypatch):
     monkeypatch.setattr(temp_dir, "_tempdir_manager", None)
     with pytest.raises(AssertionError):
         TempDirectory(globally_managed=True)
+
+
+deleted_kind = "deleted"
+not_deleted_kind = "not-deleted"
+
+
+@pytest.mark.parametrize("delete,kind,exists", [
+    (None, deleted_kind, False),
+    (True, deleted_kind, False),
+    (False, deleted_kind, True),
+    (None, not_deleted_kind, True),
+    (True, not_deleted_kind, False),
+    (False, not_deleted_kind, True),
+    (None, "unspecified", False),
+    (True, "unspecified", False),
+    (False, "unspecified", True),
+])
+def test_tempdir_registry(kind, delete, exists):
+    with tempdir_registry() as registry:
+        registry.set_delete(deleted_kind, True)
+        registry.set_delete(not_deleted_kind, False)
+
+        with TempDirectory(delete=delete, kind=kind) as d:
+            path = d.path
+            assert os.path.exists(path)
+        assert os.path.exists(path) == exists