env: delete entry in envs.toml when running poetry env remove --all

Author: radoering

src/poetry/console/commands/env/remove.py

@@ -45,5 +45,10 @@ def handle(self) -> int:
             for venv in manager.list():
                 manager.remove_venv(venv.path)
                 self.line(f"Deleted virtualenv: <comment>{venv.path}</comment>")
+            # Since we remove all the virtualenvs, we can also remove the entry
+            # in the envs file. (Strictly speaking, we should do this explicitly,
+            # in case it points to a virtualenv that had been removed manually before.)
+            if manager.envs_file.exists():
+                manager.envs_file.remove_section(manager.base_env_name)
 
         return 0

src/poetry/utils/env/env_manager.py

@@ -9,6 +9,7 @@
 import subprocess
 import sys
 
+from functools import cached_property
 from pathlib import Path
 from subprocess import CalledProcessError
 from typing import TYPE_CHECKING
@@ -45,6 +46,44 @@
     from poetry.utils.env.base_env import Env
 
 
+class EnvsFile(TOMLFile):
+    """
+    This file contains one section per project with the project's base env name
+    as section name. Each section contains the minor and patch version of the
+    python executable used to create the currently active virtualenv.
+
+    Example:
+
+    [poetry-QRErDmmj]
+    minor = "3.9"
+    patch = "3.9.13"
+
+    [poetry-core-m5r7DkRA]
+    minor = "3.11"
+    patch = "3.11.6"
+    """
+
+    def remove_section(self, name: str, minor: str | None = None) -> str | None:
+        """
+        Remove a section from the envs file.
+
+        If "minor" is given, the section is only removed if its minor value
+        matches "minor".
+
+        Returns the "minor" value of the removed section.
+        """
+        envs = self.read()
+        current_env = envs.get(name)
+        if current_env is not None and (not minor or current_env["minor"] == minor):
+            del envs[name]
+            self.write(envs)
+            minor = current_env["minor"]
+            assert isinstance(minor, str)
+            return minor
+
+        return None
+
+
 class EnvManager:
     """
     Environments manager
@@ -121,11 +160,19 @@ def in_project_venv(self) -> Path:
         venv: Path = self._poetry.file.path.parent / ".venv"
         return venv
 
+    @cached_property
+    def envs_file(self) -> EnvsFile:
+        return EnvsFile(self._poetry.config.virtualenvs_path / self.ENVS_FILE)
+
+    @cached_property
+    def base_env_name(self) -> str:
+        return self.generate_env_name(
+            self._poetry.package.name,
+            str(self._poetry.file.path.parent),
+        )
+
     def activate(self, python: str) -> Env:
         venv_path = self._poetry.config.virtualenvs_path
-        cwd = self._poetry.file.path.parent
-
-        envs_file = TOMLFile(venv_path / self.ENVS_FILE)
 
         try:
             python_version = Version.parse(python)
@@ -170,10 +217,9 @@ def activate(self, python: str) -> Env:
             return self.get(reload=True)
 
         envs = tomlkit.document()
-        base_env_name = self.generate_env_name(self._poetry.package.name, str(cwd))
-        if envs_file.exists():
-            envs = envs_file.read()
-            current_env = envs.get(base_env_name)
+        if self.envs_file.exists():
+            envs = self.envs_file.read()
+            current_env = envs.get(self.base_env_name)
             if current_env is not None:
                 current_minor = current_env["minor"]
                 current_patch = current_env["patch"]
@@ -182,7 +228,7 @@ def activate(self, python: str) -> Env:
                     # We need to recreate
                     create = True
 
-        name = f"{base_env_name}-py{minor}"
+        name = f"{self.base_env_name}-py{minor}"
         venv = venv_path / name
 
         # Create if needed
@@ -202,29 +248,21 @@ def activate(self, python: str) -> Env:
             self.create_venv(executable=python_path, force=create)
 
         # Activate
-        envs[base_env_name] = {"minor": minor, "patch": patch}
-        envs_file.write(envs)
+        envs[self.base_env_name] = {"minor": minor, "patch": patch}
+        self.envs_file.write(envs)
 
         return self.get(reload=True)
 
     def deactivate(self) -> None:
         venv_path = self._poetry.config.virtualenvs_path
-        name = self.generate_env_name(
-            self._poetry.package.name, str(self._poetry.file.path.parent)
-        )
 
-        envs_file = TOMLFile(venv_path / self.ENVS_FILE)
-        if envs_file.exists():
-            envs = envs_file.read()
-            env = envs.get(name)
-            if env is not None:
-                venv = venv_path / f"{name}-py{env['minor']}"
-                self._io.write_error_line(
-                    f"Deactivating virtualenv: <comment>{venv}</comment>"
-                )
-                del envs[name]
-
-                envs_file.write(envs)
+        if self.envs_file.exists() and (
+            minor := self.envs_file.remove_section(self.base_env_name)
+        ):
+            venv = venv_path / f"{self.base_env_name}-py{minor}"
+            self._io.write_error_line(
+                f"Deactivating virtualenv: <comment>{venv}</comment>"
+            )
 
     def get(self, reload: bool = False) -> Env:
         if self._env is not None and not reload:
@@ -237,15 +275,10 @@ def get(self, reload: bool = False) -> Env:
             precision=2, prefer_active_python=prefer_active_python, io=self._io
         ).to_string()
 
-        venv_path = self._poetry.config.virtualenvs_path
-
-        cwd = self._poetry.file.path.parent
-        envs_file = TOMLFile(venv_path / self.ENVS_FILE)
         env = None
-        base_env_name = self.generate_env_name(self._poetry.package.name, str(cwd))
-        if envs_file.exists():
-            envs = envs_file.read()
-            env = envs.get(base_env_name)
+        if self.envs_file.exists():
+            envs = self.envs_file.read()
+            env = envs.get(self.base_env_name)
             if env:
                 python_minor = env["minor"]
 
@@ -272,7 +305,7 @@ def get(self, reload: bool = False) -> Env:
 
             venv_path = self._poetry.config.virtualenvs_path
 
-            name = f"{base_env_name}-py{python_minor.strip()}"
+            name = f"{self.base_env_name}-py{python_minor.strip()}"
 
             venv = venv_path / name
 
@@ -313,12 +346,6 @@ def check_env_is_for_current_project(env: str, base_env_name: str) -> bool:
         return env.startswith(base_env_name)
 
     def remove(self, python: str) -> Env:
-        venv_path = self._poetry.config.virtualenvs_path
-
-        cwd = self._poetry.file.path.parent
-        envs_file = TOMLFile(venv_path / self.ENVS_FILE)
-        base_env_name = self.generate_env_name(self._poetry.package.name, str(cwd))
-
         python_path = Path(python)
         if python_path.is_file():
             # Validate env name if provided env is a full path to python
@@ -327,34 +354,21 @@ def remove(self, python: str) -> Env:
                     [python, "-c", GET_ENV_PATH_ONELINER], text=True
                 ).strip("\n")
                 env_name = Path(env_dir).name
-                if not self.check_env_is_for_current_project(env_name, base_env_name):
+                if not self.check_env_is_for_current_project(
+                    env_name, self.base_env_name
+                ):
                     raise IncorrectEnvError(env_name)
             except CalledProcessError as e:
                 raise EnvCommandError(e)
 
-        if self.check_env_is_for_current_project(python, base_env_name):
+        if self.check_env_is_for_current_project(python, self.base_env_name):
             venvs = self.list()
             for venv in venvs:
                 if venv.path.name == python:
                     # Exact virtualenv name
-                    if not envs_file.exists():
-                        self.remove_venv(venv.path)
-
-                        return venv
-
-                    venv_minor = ".".join(str(v) for v in venv.version_info[:2])
-                    base_env_name = self.generate_env_name(cwd.name, str(cwd))
-                    envs = envs_file.read()
-
-                    current_env = envs.get(base_env_name)
-                    if not current_env:
-                        self.remove_venv(venv.path)
-
-                        return venv
-
-                    if current_env["minor"] == venv_minor:
-                        del envs[base_env_name]
-                        envs_file.write(envs)
+                    if self.envs_file.exists():
+                        venv_minor = ".".join(str(v) for v in venv.version_info[:2])
+                        self.envs_file.remove_section(self.base_env_name, venv_minor)
 
                     self.remove_venv(venv.path)
 
@@ -389,21 +403,14 @@ def remove(self, python: str) -> Env:
         python_version = Version.parse(python_version_string.strip())
         minor = f"{python_version.major}.{python_version.minor}"
 
-        name = f"{base_env_name}-py{minor}"
+        name = f"{self.base_env_name}-py{minor}"
         venv_path = venv_path / name
 
         if not venv_path.exists():
             raise ValueError(f'<warning>Environment "{name}" does not exist.</warning>')
 
-        if envs_file.exists():
-            envs = envs_file.read()
-            current_env = envs.get(base_env_name)
-            if current_env is not None:
-                current_minor = current_env["minor"]
-
-                if current_minor == minor:
-                    del envs[base_env_name]
-                    envs_file.write(envs)
+        if self.envs_file.exists():
+            self.envs_file.remove_section(self.base_env_name, minor)
 
         self.remove_venv(venv_path)
 

tests/console/commands/env/test_remove.py

@@ -62,19 +62,50 @@ def test_remove_by_name(
     assert tester.io.fetch_output() == expected
 
 
+@pytest.mark.parametrize(
+    "envs_file", [None, "empty", "self", "other", "self_and_other"]
+)
 def test_remove_all(
     tester: CommandTester,
     venvs_in_cache_dirs: list[str],
     venv_name: str,
     venv_cache: Path,
+    envs_file: str | None,
 ) -> None:
+    envs_file_path = venv_cache / "envs.toml"
+    if envs_file == "empty":
+        envs_file_path.touch()
+    elif envs_file == "self":
+        envs_file_path.write_text(f'[{venv_name}]\nminor = "3.9"\npatch = "3.9.1"\n')
+    elif envs_file == "other":
+        envs_file_path.write_text('[other-abcdefgh]\nminor = "3.9"\npatch = "3.9.1"\n')
+    elif envs_file == "self_and_other":
+        envs_file_path.write_text(
+            f'[{venv_name}]\nminor = "3.9"\npatch = "3.9.1"\n'
+            '[other-abcdefgh]\nminor = "3.9"\npatch = "3.9.1"\n'
+        )
+    else:
+        # no envs file -> nothing to prepare
+        assert envs_file is None
+
     expected = {""}
     tester.execute("--all")
     for name in venvs_in_cache_dirs:
         assert not (venv_cache / name).exists()
         expected.add(f"Deleted virtualenv: {venv_cache / name}")
     assert set(tester.io.fetch_output().split("\n")) == expected
 
+    if envs_file is not None:
+        assert envs_file_path.exists()
+        envs_file_content = envs_file_path.read_text()
+        assert venv_name not in envs_file_content
+        if "other" in envs_file:
+            assert "other-abcdefgh" in envs_file_content
+        else:
+            assert envs_file_content == ""
+    else:
+        assert not envs_file_path.exists()
+
 
 def test_remove_all_and_version(
     tester: CommandTester,

tests/utils/env/test_env_manager.py

@@ -19,6 +19,7 @@
 from poetry.utils.env import InvalidCurrentPythonVersionError
 from poetry.utils.env import NoCompatiblePythonVersionFound
 from poetry.utils.env import PythonVersionNotFound
+from poetry.utils.env.env_manager import EnvsFile
 from poetry.utils.helpers import remove_directory
 
 
@@ -84,6 +85,41 @@ def in_project_venv_dir(poetry: Poetry) -> Iterator[Path]:
         venv_dir.rmdir()
 
 
+@pytest.mark.parametrize(
+    ("section", "version", "expected"),
+    [
+        ("foo", None, "3.10"),
+        ("bar", None, "3.11"),
+        ("baz", None, "3.12"),
+        ("bar", "3.11", "3.11"),
+        ("bar", "3.10", None),
+    ],
+)
+def test_envs_file_remove_section(
+    tmp_path: Path, section: str, version: str | None, expected: str | None
+) -> None:
+    envs_file_path = tmp_path / "envs.toml"
+
+    envs_file = TOMLFile(envs_file_path)
+    doc = tomlkit.document()
+    doc["foo"] = {"minor": "3.10", "patch": "3.10.13"}
+    doc["bar"] = {"minor": "3.11", "patch": "3.11.7"}
+    doc["baz"] = {"minor": "3.12", "patch": "3.12.1"}
+    envs_file.write(doc)
+
+    minor = EnvsFile(envs_file_path).remove_section(section, version)
+
+    assert minor == expected
+
+    envs = TOMLFile(envs_file_path).read()
+    if expected is None:
+        assert section in envs
+    else:
+        assert section not in envs
+    for other_section in {"foo", "bar", "baz"} - {section}:
+        assert other_section in envs
+
+
 def test_activate_in_project_venv_no_explicit_config(
     tmp_path: Path,
     manager: EnvManager,