Implement --prune flag for mamba env update

This fixes #688 and the similar issue for conda conda/conda#7279.

The prune flag is implemented by running the solver as if it was given
an empty prefix and then computing the diff between the current prefix
and the newly solved target and executing those resulting links and
unlinks.

This assumes that it is always possible to transform a given prefix to
any other consistent set of packages, since the solver never sees the
list of currently installed packages.

Furthermore a small fix for leaking temporary files is included.

mamba/mamba/mamba_env.py

@@ -6,6 +6,7 @@
 from __future__ import absolute_import, print_function
 
 import sys
+import os
 import tempfile
 
 from conda.base.context import context
@@ -16,7 +17,7 @@
 
 import libmambapy as api
 from mamba.linking import handle_txn
-from mamba.utils import get_installed_jsonfile, init_api_context, load_channels, to_txn
+from mamba.utils import get_installed_jsonfile, init_api_context, load_channels, to_txn, compute_final_precs, to_txn_precs
 
 
 def mamba_install(prefix, specs, args, env, dry_run=False, *_, **kwargs):
@@ -51,10 +52,21 @@ def mamba_install(prefix, specs, args, env, dry_run=False, *_, **kwargs):
 
     installed_pkg_recs = []
 
+    prune = getattr(args, 'prune', False)
+
     # We check for installed packages even while creating a new
     # Conda environment as virtual packages such as __glibc are
     # always available regardless of the environment.
     installed_json_f, installed_pkg_recs = get_installed_jsonfile(prefix)
+    if prune:
+        try:
+            installed_json_f.close()
+            os.unlink(installed_json_f.name)
+        except Exception:
+            pass
+        installed_pkg_recs_prefix = installed_pkg_recs
+        with tempfile.TemporaryDirectory() as td:
+            installed_json_f, installed_pkg_recs = get_installed_jsonfile(td)
     repo = api.Repo(pool, "installed", installed_json_f.name, "")
     repo.set_installed()
     repos.append(repo)
@@ -124,10 +136,23 @@ def mamba_install(prefix, specs, args, env, dry_run=False, *_, **kwargs):
     downloaded = transaction.prompt(repos)
     if not downloaded:
         exit(0)
-    conda_transaction = to_txn(
-        specs_to_add, [], prefix, to_link, to_unlink, installed_pkg_recs, index
-    )
+    if prune:
+        history = api.History(prefix)
+        history_map = history.get_requested_specs_map()
+        specs_to_add_names = {m.name for m in specs_to_add}
+        specs_to_remove = [MatchSpec(m) for m in history_map if m not in specs_to_add_names]
+        final_precs = compute_final_precs(None, to_link, to_unlink, installed_pkg_recs_prefix, index)
+        conda_transaction = to_txn_precs(specs_to_add, specs_to_remove, prefix, final_precs)
+    else:
+        conda_transaction = to_txn(
+            specs_to_add, [], prefix, to_link, to_unlink, installed_pkg_recs, index
+        )
     handle_txn(conda_transaction, prefix, args, True)
+    try:
+        installed_json_f.close()
+        os.unlink(installed_json_f.name)
+    except Exception:
+        pass
 
 
 def mamba_dry_run(specs, args, env, *_, **kwargs):

mamba/mamba/utils.py

@@ -336,22 +336,12 @@ def get_installed_jsonfile(prefix):
     installed_json_f.flush()
     return installed_json_f, installed_pkg_recs
 
-
-def to_txn(
-    specs_to_add,
-    specs_to_remove,
-    prefix,
-    to_link,
-    to_unlink,
-    installed_pkg_recs,
-    index=None,
-):
+def compute_final_precs(prefix_records, to_link, to_unlink, installed_pkg_recs, index=None):
     if index is None:
         index = []
     to_link_records, to_unlink_records = [], []
 
-    prefix_data = PrefixData(prefix)
-    final_precs = IndexedSet(prefix_data.iter_records())
+    final_precs = IndexedSet(prefix_records)
 
     lookup_dict = {}
     for _, entry in index:
@@ -383,9 +373,17 @@ def to_txn(
         final_precs.add(rec)
         to_link_records.append(rec)
 
+    return IndexedSet(PrefixGraph(final_precs).graph)
+
+def to_txn_precs(
+    specs_to_add,
+    specs_to_remove,
+    prefix,
+    final_precs,
+):
     unlink_precs, link_precs = diff_for_unlink_link_precs(
         prefix,
-        final_precs=IndexedSet(PrefixGraph(final_precs).graph),
+        final_precs=final_precs,
         specs_to_add=specs_to_add,
         force_reinstall=context.force_reinstall,
     )
@@ -401,3 +399,27 @@ def to_txn(
 
     conda_transaction = UnlinkLinkTransaction(pref_setup)
     return conda_transaction
+
+def to_txn(
+    specs_to_add,
+    specs_to_remove,
+    prefix,
+    to_link,
+    to_unlink,
+    installed_pkg_recs,
+    index=None,
+):
+    prefix_data = PrefixData(prefix)
+    final_precs = compute_final_precs(
+        prefix_data.iter_records(),
+        to_link,
+        to_unlink,
+        installed_pkg_recs,
+        index,
+    )
+    return to_txn_precs(
+        specs_to_add,
+        specs_to_remove,
+        prefix,
+        final_precs,
+    )