[micromamba] Fix behavior of env update (to mimick conda) (#3396)

* Fix behavior of env update (to mimick conda) * Code review: Use struct and enum class for update parameters
mamba-org · Aug 8, 2024 · d17ebc5 · d17ebc5
1 parent c7563b8
commit d17ebc5
Show file tree

Hide file tree

Showing 6 changed files with 136 additions and 52 deletions.
diff --git a/libmamba/include/mamba/api/update.hpp b/libmamba/include/mamba/api/update.hpp
@@ -11,12 +11,39 @@ namespace mamba
 {
     class Configuration;
 
-    void update(
-        Configuration& config,
-        bool update_all = false,
-        bool prune_deps = false,
-        bool remove_not_specified = false
-    );
+    enum class UpdateAll : bool
+    {
+        No = false,
+        Yes = true,
+    };
+
+    enum class PruneDeps : bool
+    {
+        No = false,
+        Yes = true,
+    };
+
+    enum class EnvUpdate : bool  // Specific to `env update` command
+    {
+        No = false,
+        Yes = true,
+    };
+
+    enum class RemoveNotSpecified : bool  // Specific to `env update` command
+    {
+        No = false,
+        Yes = true,
+    };
+
+    struct UpdateParams
+    {
+        UpdateAll update_all = UpdateAll::No;
+        PruneDeps prune_deps = PruneDeps::No;
+        EnvUpdate env_update = EnvUpdate::No;
+        RemoveNotSpecified remove_not_specified = RemoveNotSpecified::No;
+    };
+
+    void update(Configuration& config, const UpdateParams& update_params = {});
 }
 
 #endif
diff --git a/libmamba/src/api/c_api.cpp b/libmamba/src/api/c_api.cpp
@@ -101,7 +101,9 @@ mamba_update(mamba::Configuration* config, int update_all)
     assert(config != nullptr);
     try
     {
-        update(*config, update_all);
+        mamba::UpdateParams update_params{};
+        update_params.update_all = update_all ? mamba::UpdateAll::Yes : mamba::UpdateAll::No;
+        update(*config, update_params);
         return 0;
     }
     catch (...)

diff --git a/libmamba/src/api/update.cpp b/libmamba/src/api/update.cpp
@@ -25,18 +25,16 @@ namespace mamba
         auto create_update_request(
             PrefixData& prefix_data,
             std::vector<std::string> specs,
-            bool update_all,
-            bool prune_deps,
-            bool remove_not_specified
+            const UpdateParams& update_params
         ) -> solver::Request
         {
             using Request = solver::Request;
 
             auto request = Request();
 
-            if (update_all)
+            if (update_params.update_all == UpdateAll::Yes)
             {
-                if (prune_deps)
+                if (update_params.prune_deps == PruneDeps::Yes)
                 {
                     auto hist_map = prefix_data.history().get_requested_specs_map();
                     request.jobs.reserve(hist_map.size() + 1);
@@ -55,39 +53,77 @@ namespace mamba
             else
             {
                 request.jobs.reserve(specs.size());
-                if (remove_not_specified)
+                if (update_params.env_update == EnvUpdate::Yes)
                 {
-                    auto hist_map = prefix_data.history().get_requested_specs_map();
-                    for (auto& it : hist_map)
+                    if (update_params.remove_not_specified == RemoveNotSpecified::Yes)
                     {
-                        if (std::find(specs.begin(), specs.end(), it.second.name().str())
-                            == specs.end())
+                        auto hist_map = prefix_data.history().get_requested_specs_map();
+                        for (auto& it : hist_map)
                         {
-                            request.jobs.emplace_back(Request::Remove{
-                                specs::MatchSpec::parse(it.second.name().str())
-                                    .or_else([](specs::ParseError&& err) { throw std::move(err); })
-                                    .value(),
-                                /* .clean_dependencies= */ true,
-                            });
+                            // We use `spec_names` here because `specs` contain more info than just
+                            // the spec name.
+                            // Therefore, the search later and comparison (using `specs`) with
+                            // MatchSpec.name().str() in `hist_map` second elements wouldn't be
+                            // relevant
+                            std::vector<std::string> spec_names;
+                            spec_names.reserve(specs.size());
+                            std::transform(
+                                specs.begin(),
+                                specs.end(),
+                                std::back_inserter(spec_names),
+                                [](const std::string& spec)
+                                {
+                                    return specs::MatchSpec::parse(spec)
+                                        .or_else([](specs::ParseError&& err)
+                                                 { throw std::move(err); })
+                                        .value()
+                                        .name()
+                                        .str();
+                                }
+                            );
+
+                            if (std::find(spec_names.begin(), spec_names.end(), it.second.name().str())
+                                == spec_names.end())
+                            {
+                                request.jobs.emplace_back(Request::Remove{
+                                    specs::MatchSpec::parse(it.second.name().str())
+                                        .or_else([](specs::ParseError&& err)
+                                                 { throw std::move(err); })
+                                        .value(),
+                                    /* .clean_dependencies= */ true,
+                                });
+                            }
                         }
                     }
-                }
 
-                for (const auto& raw_ms : specs)
+                    // Install/update everything in specs
+                    for (const auto& raw_ms : specs)
+                    {
+                        request.jobs.emplace_back(Request::Install{
+                            specs::MatchSpec::parse(raw_ms)
+                                .or_else([](specs::ParseError&& err) { throw std::move(err); })
+                                .value(),
+                        });
+                    }
+                }
+                else
                 {
-                    request.jobs.emplace_back(Request::Update{
-                        specs::MatchSpec::parse(raw_ms)
-                            .or_else([](specs::ParseError&& err) { throw std::move(err); })
-                            .value(),
-                    });
+                    for (const auto& raw_ms : specs)
+                    {
+                        request.jobs.emplace_back(Request::Update{
+                            specs::MatchSpec::parse(raw_ms)
+                                .or_else([](specs::ParseError&& err) { throw std::move(err); })
+                                .value(),
+                        });
+                    }
                 }
             }
 
             return request;
         }
     }
 
-    void update(Configuration& config, bool update_all, bool prune_deps, bool remove_not_specified)
+    void update(Configuration& config, const UpdateParams& update_params)
     {
         auto& ctx = config.context();
 
@@ -139,13 +175,7 @@ namespace mamba
 
         load_installed_packages_in_database(ctx, db, prefix_data);
 
-        auto request = create_update_request(
-            prefix_data,
-            raw_update_specs,
-            /* update_all= */ update_all,
-            /* prune_deps= */ prune_deps,
-            /* remove_not_specified= */ remove_not_specified
-        );
+        auto request = create_update_request(prefix_data, raw_update_specs, update_params);
         add_pins_to_request(
             request,
             ctx,
@@ -187,17 +217,17 @@ namespace mamba
         );
 
 
-        auto execute_transaction = [&](MTransaction& transaction)
+        auto execute_transaction = [&](MTransaction& trans)
         {
             if (ctx.output_params.json)
             {
-                transaction.log_json();
+                trans.log_json();
             }
 
-            bool yes = transaction.prompt(ctx, channel_context);
+            bool yes = trans.prompt(ctx, channel_context);
             if (yes)
             {
-                transaction.execute(ctx, channel_context, prefix_data);
+                trans.execute(ctx, channel_context, prefix_data);
             }
         };
 

diff --git a/micromamba/src/env.cpp b/micromamba/src/env.cpp
@@ -284,6 +284,15 @@ set_env_command(CLI::App* com, Configuration& config)
 
     update_subcom->callback(
         [&config]
-        { update(config, /*update_all*/ false, /*prune_deps*/ false, remove_not_specified); }
+        {
+            auto update_params = UpdateParams{
+                UpdateAll::No,
+                PruneDeps::Yes,
+                EnvUpdate::Yes,
+                remove_not_specified ? RemoveNotSpecified::Yes : RemoveNotSpecified::No,
+            };
+
+            update(config, update_params);
+        }
     );
 }
diff --git a/micromamba/src/update.cpp b/micromamba/src/update.cpp
@@ -209,7 +209,18 @@ set_update_command(CLI::App* subcom, Configuration& config)
     subcom->get_option("specs")->description("Specs to update in the environment");
     subcom->add_flag("-a,--all", update_all, "Update all packages in the environment");
 
-    subcom->callback([&] { return update(config, update_all, prune_deps); });
+    subcom->callback(
+        [&]
+        {
+            auto update_params = UpdateParams{
+                update_all ? UpdateAll::Yes : UpdateAll::No,
+                prune_deps ? PruneDeps::Yes : PruneDeps::No,
+                EnvUpdate::No,
+                RemoveNotSpecified::No,
+            };
+            return update(config, update_params);
+        }
+    );
 }
 
 void

diff --git a/micromamba/tests/test_env.py b/micromamba/tests/test_env.py
@@ -125,11 +125,12 @@ def test_env_remove(tmp_home, tmp_root_prefix):
         assert str(env_fp) not in lines
 
 
-env_yaml_content = """
+env_yaml_content_with_version_and_new_pkg = """
 channels:
 - conda-forge
 dependencies:
-- python
+- python 3.11.*
+- ipython
 """
 
 
@@ -138,35 +139,39 @@ def test_env_remove(tmp_home, tmp_root_prefix):
 def test_env_update(tmp_home, tmp_root_prefix, tmp_path, prune):
     env_name = "env-create-update"
 
-    # Create env with python=3.11.0 and xtensor=0.25.0
-    helpers.create("python=3.11.0", "xtensor=0.25.0", "-n", env_name, "--json", no_dry_run=True)
+    # Create env with python=3.10.0 and xtensor=0.25.0
+    helpers.create("python=3.10.0", "xtensor=0.25.0", "-n", env_name, "--json", no_dry_run=True)
     packages = helpers.umamba_list("-n", env_name, "--json")
     assert any(
-        package["name"] == "python" and package["version"] == "3.11.0" for package in packages
+        package["name"] == "python" and package["version"] == "3.10.0" for package in packages
     )
     assert any(
         package["name"] == "xtensor" and package["version"] == "0.25.0" for package in packages
     )
     assert any(package["name"] == "xtl" for package in packages)
+    assert not any(package["name"] == "ipython" for package in packages)
 
     # Update python
     from packaging.version import Version
 
     env_file_yml = tmp_path / "test_env.yaml"
-    env_file_yml.write_text(env_yaml_content)
+    env_file_yml.write_text(env_yaml_content_with_version_and_new_pkg)
 
     cmd = ["update", "-n", env_name, f"--file={env_file_yml}", "-y"]
     if prune:
         cmd += ["--prune"]
     helpers.run_env(*cmd)
     packages = helpers.umamba_list("-n", env_name, "--json")
     assert any(
-        package["name"] == "python" and Version(package["version"]) > Version("3.11.0")
+        package["name"] == "python" and Version(package["version"]) >= Version("3.11.0")
         for package in packages
     )
+    assert any(package["name"] == "ipython" for package in packages)
     if prune:
         assert not any(package["name"] == "xtensor" for package in packages)
-        # Make sure dependencies of removed pkgs are removed as well (xtl is a dep of xtensor)
+        # Make sure dependencies of removed pkgs are removed as well
+        # since `prune_deps` is always set to true in the case of `env update` command
+        # (xtl is a dep of xtensor)
         assert not any(package["name"] == "xtl" for package in packages)
     else:
         assert any(