Make mamba env download and extract using libmamba

mamba/mamba/linking.py

@@ -0,0 +1,50 @@
+# Copyright (C) 2019, QuantStack
+# SPDX-License-Identifier: BSD-3-Clause
+
+from conda.base.context import context
+from conda.cli import common as cli_common
+from conda.core.package_cache_data import PackageCacheData
+from conda.exceptions import (
+    CondaExitZero,
+    CondaSystemExit,
+    DryRunExit,
+    PackagesNotFoundError,
+)
+
+from mamba.utils import lock_file
+
+
+def handle_txn(unlink_link_transaction, prefix, args, newenv, remove_op=False):
+    if unlink_link_transaction.nothing_to_do:
+        if remove_op:
+            # No packages found to remove from environment
+            raise PackagesNotFoundError(args.package_names)
+        elif not newenv:
+            if context.json:
+                cli_common.stdout_json_success(
+                    message="All requested packages already installed."
+                )
+            return
+
+    if context.dry_run:
+        actions = unlink_link_transaction._make_legacy_action_groups()[0]
+        if context.json:
+            cli_common.stdout_json_success(prefix=prefix, actions=actions, dry_run=True)
+        raise DryRunExit()
+
+    try:
+        with lock_file(PackageCacheData.first_writable().pkgs_dir):
+            unlink_link_transaction.download_and_extract()
+            if context.download_only:
+                raise CondaExitZero(
+                    "Package caches prepared. UnlinkLinkTransaction cancelled with "
+                    "--download-only option."
+                )
+            unlink_link_transaction.execute()
+
+    except SystemExit as e:
+        raise CondaSystemExit("Exiting", e)
+
+    if context.json:
+        actions = unlink_link_transaction._make_legacy_action_groups()[0]
+        cli_common.stdout_json_success(prefix=prefix, actions=actions)

mamba/mamba/mamba.py

@@ -12,7 +12,6 @@
 # create support
 from conda.base.constants import ChannelPriority, DepsModifier, UpdateModifier
 from conda.base.context import context
-from conda.cli import common as cli_common
 from conda.cli.common import (
     check_non_admin,
     confirm_yn,
@@ -31,12 +30,9 @@
 from conda.exceptions import (
     ArgumentError,
     CondaEnvironmentError,
-    CondaExitZero,
     CondaOSError,
-    CondaSystemExit,
     CondaValueError,
     DirectoryNotACondaEnvironmentError,
-    DryRunExit,
     EnvironmentLocationNotFound,
     NoBaseEnvironmentError,
     PackageNotInstalledError,
@@ -52,6 +48,7 @@
 import libmambapy as api
 import mamba
 from mamba import repoquery as repoquery_api
+from mamba.linking import handle_txn
 from mamba.mamba_shell_init import shell_init
 from mamba.utils import (
     get_installed_jsonfile,
@@ -128,46 +125,6 @@ def arg2spec(arg, json=False, update=False):
 use_mamba_experimental = False
 
 
-def handle_txn(unlink_link_transaction, prefix, args, newenv, remove_op=False):
-    if unlink_link_transaction.nothing_to_do:
-        if remove_op:
-            # No packages found to remove from environment
-            raise PackagesNotFoundError(args.package_names)
-        elif not newenv:
-            if context.json:
-                cli_common.stdout_json_success(
-                    message="All requested packages already installed."
-                )
-            return
-
-    if context.dry_run:
-        actions = unlink_link_transaction._make_legacy_action_groups()[0]
-        if context.json:
-            cli_common.stdout_json_success(prefix=prefix, actions=actions, dry_run=True)
-        raise DryRunExit()
-
-    try:
-        with lock_file(PackageCacheData.first_writable().pkgs_dir):
-            unlink_link_transaction.download_and_extract()
-            if context.download_only:
-                raise CondaExitZero(
-                    "Package caches prepared. UnlinkLinkTransaction cancelled with "
-                    "--download-only option."
-                )
-            unlink_link_transaction.execute()
-
-    except SystemExit as e:
-        raise CondaSystemExit("Exiting", e)
-
-    if newenv:
-        touch_nonadmin(prefix)
-        print_activate(args.name if args.name else prefix)
-
-    if context.json:
-        actions = unlink_link_transaction._make_legacy_action_groups()[0]
-        cli_common.stdout_json_success(prefix=prefix, actions=actions)
-
-
 def remove(args, parser):
     if not (args.all or args.package_names):
         raise CondaValueError(
@@ -636,6 +593,10 @@ def install(args, parser, command="install"):
         )
         handle_txn(conda_transaction, prefix, args, newenv)
 
+    if newenv:
+        touch_nonadmin(prefix)
+        print_activate(args.name if args.name else prefix)
+
     try:
         installed_json_f.close()
         os.unlink(installed_json_f.name)

mamba/mamba/mamba_env.py

@@ -14,6 +14,7 @@
 from conda_env.installers import conda
 
 import libmambapy as api
+from mamba.linking import handle_txn
 from mamba.utils import get_installed_jsonfile, init_api_context, load_channels, to_txn
 
 
@@ -118,13 +119,15 @@ def mamba_install(prefix, specs, args, env, *_, **kwargs):
 
     specs_to_add = [MatchSpec(m) for m in mmb_specs[0]]
 
+    transaction.log_json()
+    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
     )
-
-    pfe = conda_transaction._get_pfe()
-    pfe.execute()
-    conda_transaction.execute()
+    handle_txn(conda_transaction, prefix, args, True)
 
 
 conda.install = mamba_install