From 089ff9dd8adbb6384ad007eee7e81c89d2a1cd68 Mon Sep 17 00:00:00 2001
From: DustinMoriarty <dustin.moriarty@protonmail.com>
Date: Sun, 14 Nov 2021 17:45:12 -0600
Subject: [PATCH] Allow removing multiple (and all) environments at once
 (#3212)

Co-authored-by: Arun Babu Neelicattu <arun.neelicattu@gmail.com>
Co-authored-by: Bjorn Neergaard <bjorn@neersighted.com>
---
 docs/managing-environments.md             |  9 ++++++
 poetry/console/commands/env/remove.py     | 34 +++++++++++++++++++----
 tests/console/commands/env/test_remove.py | 31 +++++++++++++++++++++
 3 files changed, 69 insertions(+), 5 deletions(-)

diff --git a/docs/managing-environments.md b/docs/managing-environments.md
index 5f2a98d8059..7a4b21a6643 100644
--- a/docs/managing-environments.md
+++ b/docs/managing-environments.md
@@ -129,4 +129,13 @@ poetry env remove 3.7
 poetry env remove test-O3eWbxRl-py3.7
 ```
 
+You can delete more than one environment at a time.
+```bash
+poetry env remove python3.6 python3.7 python3.8
+```
+Use the `--all` option to delete all virtual environments at once.
+```bash
+poetry env remove --all
+```
+
 If you remove the currently activated virtual environment, it will be automatically deactivated.
diff --git a/poetry/console/commands/env/remove.py b/poetry/console/commands/env/remove.py
index 9d5153c420b..8b95ce486dd 100644
--- a/poetry/console/commands/env/remove.py
+++ b/poetry/console/commands/env/remove.py
@@ -1,4 +1,5 @@
 from cleo.helpers import argument
+from cleo.helpers import option
 
 from ..command import Command
 
@@ -6,16 +7,39 @@
 class EnvRemoveCommand(Command):
 
     name = "env remove"
-    description = "Removes a specific virtualenv associated with the project."
+    description = "Remove virtual environments associated with the project."
 
     arguments = [
-        argument("python", "The python executable to remove the virtualenv for.")
+        argument(
+            "python",
+            "The python executables associated with, or names of the virtual environments which are to "
+            "be removed.",
+            optional=True,
+            multiple=True,
+        )
+    ]
+    options = [
+        option(
+            "all",
+            description="Remove all managed virtual environments associated with the "
+            "project.",
+        ),
     ]
 
     def handle(self) -> None:
         from poetry.utils.env import EnvManager
 
-        manager = EnvManager(self.poetry)
-        venv = manager.remove(self.argument("python"))
+        pythons = self.argument("python")
+        all = self.option("all")
+        if not (pythons or all):
+            self.line("No virtualenv provided.")
 
-        self.line("Deleted virtualenv: <comment>{}</comment>".format(venv.path))
+        manager = EnvManager(self.poetry)
+        # TODO: refactor env.py to allow removal with one loop
+        for python in pythons:
+            venv = manager.remove(python)
+            self.line("Deleted virtualenv: <comment>{}</comment>".format(venv.path))
+        if all:
+            for venv in manager.list():
+                manager.remove_venv(venv.path)
+                self.line("Deleted virtualenv: <comment>{}</comment>".format(venv.path))
diff --git a/tests/console/commands/env/test_remove.py b/tests/console/commands/env/test_remove.py
index b95419bb6fb..d79bdca3ec5 100644
--- a/tests/console/commands/env/test_remove.py
+++ b/tests/console/commands/env/test_remove.py
@@ -39,3 +39,34 @@ def test_remove_by_name(tester, venvs_in_cache_dirs, venv_name, venv_cache):
         expected += "Deleted virtualenv: {}\n".format((venv_cache / name))
 
     assert expected == tester.io.fetch_output()
+
+
+def test_remove_all(tester, venvs_in_cache_dirs, venv_name, venv_cache):
+    expected = {""}
+    tester.execute("--all")
+    for name in venvs_in_cache_dirs:
+        assert not (venv_cache / name).exists()
+        expected.add("Deleted virtualenv: {}".format((venv_cache / name)))
+    assert expected == set(tester.io.fetch_output().split("\n"))
+
+
+def test_remove_all_and_version(tester, venvs_in_cache_dirs, venv_name, venv_cache):
+    expected = {""}
+    tester.execute("--all {}".format(venvs_in_cache_dirs[0]))
+    for name in venvs_in_cache_dirs:
+        assert not (venv_cache / name).exists()
+        expected.add("Deleted virtualenv: {}".format((venv_cache / name)))
+    assert expected == set(tester.io.fetch_output().split("\n"))
+
+
+def test_remove_multiple(tester, venvs_in_cache_dirs, venv_name, venv_cache):
+    expected = {""}
+    removed_envs = venvs_in_cache_dirs[0:2]
+    remaining_envs = venvs_in_cache_dirs[2:]
+    tester.execute(" ".join(removed_envs))
+    for name in removed_envs:
+        assert not (venv_cache / name).exists()
+        expected.add("Deleted virtualenv: {}".format((venv_cache / name)))
+    for name in remaining_envs:
+        assert (venv_cache / name).exists()
+    assert expected == set(tester.io.fetch_output().split("\n"))