Skip to content

Commit 0a57e4e

Browse files
chrahuntxavfernandez
authored andcommitted
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).
1 parent e0b2d96 commit 0a57e4e

File tree

2 files changed

+73
-2
lines changed

2 files changed

+73
-2
lines changed

src/pip/_internal/utils/temp_dir.py

+46-2
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
from pip._internal.utils.typing import MYPY_CHECK_RUNNING
1414

1515
if MYPY_CHECK_RUNNING:
16-
from typing import Any, Iterator, Optional, TypeVar
16+
from typing import Any, Dict, Iterator, Optional, TypeVar
1717

1818
_T = TypeVar('_T', bound='TempDirectory')
1919

@@ -36,6 +36,47 @@ def global_tempdir_manager():
3636
_tempdir_manager = old_tempdir_manager
3737

3838

39+
class TempDirectoryTypeRegistry(object):
40+
"""Manages temp directory behavior
41+
"""
42+
43+
def __init__(self):
44+
# type: () -> None
45+
self._should_delete = {} # type: Dict[str, bool]
46+
47+
def set_delete(self, kind, value):
48+
# type: (str, bool) -> None
49+
"""Indicate whether a TempDirectory of the given kind should be
50+
auto-deleted.
51+
"""
52+
self._should_delete[kind] = value
53+
54+
def get_delete(self, kind):
55+
# type: (str) -> bool
56+
"""Get configured auto-delete flag for a given TempDirectory type,
57+
default True.
58+
"""
59+
return self._should_delete.get(kind, True)
60+
61+
62+
_tempdir_registry = None # type: Optional[TempDirectoryTypeRegistry]
63+
64+
65+
@contextmanager
66+
def tempdir_registry():
67+
# type: () -> Iterator[TempDirectoryTypeRegistry]
68+
"""Provides a scoped global tempdir registry that can be used to dictate
69+
whether directories should be deleted.
70+
"""
71+
global _tempdir_registry
72+
old_tempdir_registry = _tempdir_registry
73+
_tempdir_registry = TempDirectoryTypeRegistry()
74+
try:
75+
yield _tempdir_registry
76+
finally:
77+
_tempdir_registry = old_tempdir_registry
78+
79+
3980
class TempDirectory(object):
4081
"""Helper class that owns and cleans up a temporary directory.
4182
@@ -68,8 +109,11 @@ def __init__(
68109

69110
if path is None and delete is None:
70111
# If we were not given an explicit directory, and we were not given
71-
# an explicit delete option, then we'll default to deleting.
112+
# an explicit delete option, then we'll default to deleting unless
113+
# the tempdir_registry says otherwise.
72114
delete = True
115+
if _tempdir_registry:
116+
delete = _tempdir_registry.get_delete(kind)
73117

74118
if path is None:
75119
path = self._create(kind)

tests/unit/test_utils_temp_dir.py

+27
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
AdjacentTempDirectory,
1212
TempDirectory,
1313
global_tempdir_manager,
14+
tempdir_registry,
1415
)
1516

1617

@@ -207,3 +208,29 @@ def test_tempdirectory_asserts_global_tempdir(monkeypatch):
207208
monkeypatch.setattr(temp_dir, "_tempdir_manager", None)
208209
with pytest.raises(AssertionError):
209210
TempDirectory(globally_managed=True)
211+
212+
213+
deleted_kind = "deleted"
214+
not_deleted_kind = "not-deleted"
215+
216+
217+
@pytest.mark.parametrize("delete,kind,exists", [
218+
(None, deleted_kind, False),
219+
(True, deleted_kind, False),
220+
(False, deleted_kind, True),
221+
(None, not_deleted_kind, True),
222+
(True, not_deleted_kind, False),
223+
(False, not_deleted_kind, True),
224+
(None, "unspecified", False),
225+
(True, "unspecified", False),
226+
(False, "unspecified", True),
227+
])
228+
def test_tempdir_registry(kind, delete, exists):
229+
with tempdir_registry() as registry:
230+
registry.set_delete(deleted_kind, True)
231+
registry.set_delete(not_deleted_kind, False)
232+
233+
with TempDirectory(delete=delete, kind=kind) as d:
234+
path = d.path
235+
assert os.path.exists(path)
236+
assert os.path.exists(path) == exists

0 commit comments

Comments
 (0)