Merge pull request #7652 from chrahunt/refactor/configure-tempdir-registry-in-command

Configure tempdir registry in BaseCommand

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())

src/pip/_internal/utils/temp_dir.py

@@ -13,7 +13,7 @@
 from pip._internal.utils.typing import MYPY_CHECK_RUNNING
 
 if MYPY_CHECK_RUNNING:
-    from typing import Any, Dict, Iterator, Optional, TypeVar
+    from typing import Any, Dict, Iterator, Optional, TypeVar, Union
 
     _T = TypeVar('_T', bound='TempDirectory')
 
@@ -77,6 +77,13 @@ def tempdir_registry():
         _tempdir_registry = old_tempdir_registry
 
 
+class _Default(object):
+    pass
+
+
+_default = _Default()
+
+
 class TempDirectory(object):
     """Helper class that owns and cleans up a temporary directory.
 
@@ -101,16 +108,21 @@ class TempDirectory(object):
     def __init__(
         self,
         path=None,    # type: Optional[str]
-        delete=None,  # type: Optional[bool]
+        delete=_default,  # type: Union[bool, None, _Default]
         kind="temp",  # type: str
         globally_managed=False,  # type: bool
     ):
         super(TempDirectory, self).__init__()
 
-        # If we were given an explicit directory, resolve delete option now.
-        # Otherwise we wait until cleanup and see what tempdir_registry says.
-        if path is not None and delete is None:
-            delete = False
+        if delete is _default:
+            if path is not None:
+                # If we were given an explicit directory, resolve delete option
+                # now.
+                delete = False
+            else:
+                # Otherwise, we wait until cleanup and see what
+                # tempdir_registry says.
+                delete = None
 
         if path is None:
             path = self._create(kind)

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()

tests/unit/test_utils_temp_dir.py

@@ -10,6 +10,7 @@
 from pip._internal.utils.temp_dir import (
     AdjacentTempDirectory,
     TempDirectory,
+    _default,
     global_tempdir_manager,
     tempdir_registry,
 )
@@ -216,12 +217,15 @@ def test_tempdirectory_asserts_global_tempdir(monkeypatch):
 
 @pytest.mark.parametrize("delete,kind,exists", [
     (None, deleted_kind, False),
+    (_default, deleted_kind, False),
     (True, deleted_kind, False),
     (False, deleted_kind, True),
     (None, not_deleted_kind, True),
+    (_default, not_deleted_kind, True),
     (True, not_deleted_kind, False),
     (False, not_deleted_kind, True),
     (None, "unspecified", False),
+    (_default, "unspecified", False),
     (True, "unspecified", False),
     (False, "unspecified", True),
 ])
@@ -236,6 +240,24 @@ def test_tempdir_registry(kind, delete, exists):
         assert os.path.exists(path) == exists
 
 
+@pytest.mark.parametrize("delete,exists", [
+    (_default, True), (None, False)
+])
+def test_temp_dir_does_not_delete_explicit_paths_by_default(
+    tmpdir, delete, exists
+):
+    path = tmpdir / "example"
+    path.mkdir()
+
+    with tempdir_registry() as registry:
+        registry.set_delete(deleted_kind, True)
+
+        with TempDirectory(path=path, delete=delete, kind=deleted_kind) as d:
+            assert str(d.path) == path
+            assert os.path.exists(path)
+        assert os.path.exists(path) == exists
+
+
 @pytest.mark.parametrize("should_delete", [True, False])
 def test_tempdir_registry_lazy(should_delete):
     """