Handle environment variables from `yaml` file (#3955)

libmamba/include/mamba/api/install.hpp

@@ -8,6 +8,7 @@
 #define MAMBA_API_INSTALL_HPP
 
 #include <iosfwd>
+#include <map>
 #include <string>
 #include <vector>
 
@@ -80,7 +81,18 @@ namespace mamba
     {
         void create_target_directory(const Context& context, const fs::u8path prefix);
 
-        void create_empty_target(const Context& context, const fs::u8path& prefix);
+        void create_empty_target(
+            const Context& context,
+            const fs::u8path& prefix,
+            const std::map<std::string, std::string>& env_vars,
+            bool no_env
+        );
+
+        void populate_state_file(
+            const fs::u8path& prefix,
+            const std::map<std::string, std::string>& env_vars,
+            bool no_env
+        );
 
         void file_specs_hook(Configuration& config, std::vector<std::string>& file_specs);
 
@@ -102,6 +114,7 @@ namespace mamba
         {
             std::string name;
             std::vector<std::string> dependencies, channels;
+            std::map<std::string, std::string> variables;
             std::vector<other_pkg_mgr_spec> others_pkg_mgrs_specs;
         };
 

libmamba/src/api/configuration.cpp

@@ -1395,6 +1395,12 @@ namespace mamba
                    .set_post_merge_hook(detail::file_spec_env_name_hook)
                    .description("Name of the target prefix, specified in a YAML spec file"));
 
+        insert(Configurable("spec_file_env_vars", std::map<std::string, std::string>({}))
+                   .group("Basic")
+                   .needs({ "file_specs" })
+                   .set_single_op_lifetime()
+                   .description("Environment variables specified in a YAML spec file"));
+
         insert(Configurable("specs", std::vector<std::string>({}))
                    .group("Basic")
                    .needs({ "file_specs" })  // explicit file specs overwrite current specs

libmamba/src/api/create.cpp

@@ -32,6 +32,8 @@ namespace mamba
         auto& create_specs = config.at("specs").value<std::vector<std::string>>();
         auto& use_explicit = config.at("explicit_install").value<bool>();
         auto& json_format = config.at("json").get_cli_config<bool>();
+        auto& env_vars = config.at("spec_file_env_vars").value<std::map<std::string, std::string>>();
+        auto& no_env = config.at("no_env").value<bool>();
 
         auto channel_context = ChannelContext::make_conda_compatible(ctx);
 
@@ -82,7 +84,7 @@ namespace mamba
             }
             if (create_specs.empty())
             {
-                detail::create_empty_target(ctx, ctx.prefix_params.target_prefix);
+                detail::create_empty_target(ctx, ctx.prefix_params.target_prefix, env_vars, no_env);
             }
 
             if (config.at("platform").configured() && !config.at("platform").rc_configured())

libmamba/src/api/install.cpp

@@ -281,9 +281,20 @@ namespace mamba
             {
                 result.name = f["name"].as<std::string>();
             }
+            else
             {
                 LOG_DEBUG << "No env 'name' specified in YAML spec file '" << file.string() << "'";
             }
+
+            if (f["variables"])
+            {
+                result.variables = f["variables"].as<std::map<std::string, std::string>>();
+            }
+            else
+            {
+                LOG_DEBUG << "No 'variables' specified in YAML spec file '" << file.string() << "'";
+            }
+
             return result;
         }
 
@@ -498,6 +509,8 @@ namespace mamba
             auto& no_py_pin = config.at("no_py_pin").value<bool>();
             auto& freeze_installed = config.at("freeze_installed").value<bool>();
             auto& retry_clean_cache = config.at("retry_clean_cache").value<bool>();
+            auto& env_vars = config.at("spec_file_env_vars").value<std::map<std::string, std::string>>();
+            auto& no_env = config.at("no_env").value<bool>();
 
             if (ctx.prefix_params.target_prefix.empty())
             {
@@ -647,6 +660,8 @@ namespace mamba
                     detail::create_target_directory(ctx, ctx.prefix_params.target_prefix);
                 }
 
+                detail::populate_state_file(ctx.prefix_params.target_prefix, env_vars, no_env);
+
                 trans.execute(ctx, channel_context, prefix_data);
 
                 // Print activation message only if the environment is freshly created
@@ -860,10 +875,17 @@ namespace mamba
             other
         };
 
-        void create_empty_target(const Context& context, const fs::u8path& prefix)
+        void create_empty_target(
+            const Context& context,
+            const fs::u8path& prefix,
+            const std::map<std::string, std::string>& env_vars,
+            bool no_env
+        )
         {
             detail::create_target_directory(context, prefix);
 
+            populate_state_file(prefix, env_vars, no_env);
+
             Console::instance().print(util::join(
                 "",
                 std::vector<std::string>({ "Empty environment created at prefix: ", prefix.string() })
@@ -871,6 +893,41 @@ namespace mamba
             Console::instance().json_write({ { "success", true } });
         }
 
+        void populate_state_file(
+            const fs::u8path& prefix,
+            const std::map<std::string, std::string>& env_vars,
+            bool no_env
+        )
+        {
+            if (!env_vars.empty())
+            {
+                if (!no_env)
+                {
+                    fs::u8path env_vars_file_path = prefix / "conda-meta" / "state";
+
+                    if (!fs::exists(env_vars_file_path))
+                    {
+                        path::touch(env_vars_file_path, true);
+                    }
+                    std::ofstream out = open_ofstream(env_vars_file_path, std::ios::app);
+                    if (out.fail())
+                    {
+                        throw std::runtime_error("Couldn't open file: " + env_vars_file_path.string());
+                    }
+                    else
+                    {
+                        nlohmann::json j;
+                        j["env_vars"] = env_vars;
+                        out << j.dump();
+                    }
+                }
+                else
+                {
+                    LOG_WARNING << "Using `no-env`. Variables from yaml file are not considered.";
+                }
+            }
+        }
+
         void create_target_directory(const Context& context, const fs::u8path prefix)
         {
             path::touch(prefix / "conda-meta" / "history", true);
@@ -886,6 +943,7 @@ namespace mamba
             auto& specs = config.at("specs");
             auto& others_pkg_mgrs_specs = config.at("others_pkg_mgrs_specs");
             auto& channels = config.at("channels");
+            auto& env_vars = config.at("spec_file_env_vars");
 
             auto& context = config.context();
 
@@ -999,6 +1057,20 @@ namespace mamba
                         }
                         others_pkg_mgrs_specs.set_cli_value(updated_specs);
                     }
+
+                    if (parse_result.variables.size() != 0)
+                    {
+                        std::map<std::string, std::string> updated_env_vars;
+                        if (env_vars.cli_configured())
+                        {
+                            updated_env_vars = env_vars.cli_value<std::map<std::string, std::string>>();
+                        }
+                        updated_env_vars.insert(
+                            parse_result.variables.cbegin(),
+                            parse_result.variables.cend()
+                        );
+                        env_vars.set_cli_value(updated_env_vars);
+                    }
                 }
                 else
                 {

micromamba/tests/test_create.py

@@ -419,6 +419,49 @@ def test_channels(tmp_home, tmp_root_prefix, tmp_path, cli, yaml, env_var, rc_fi
         assert res["channels"] == ["conda-forge"]
 
 
+@pytest.mark.parametrize("shared_pkgs_dirs", [True], indirect=True)
+@pytest.mark.parametrize("env_vars", (False, True))
+@pytest.mark.parametrize("no_env", (False, True))
+def test_spec_file_env_vars(tmp_home, tmp_root_prefix, tmp_path, env_vars, no_env):
+    env_name = "env-check-env-vars"
+
+    spec_file = tmp_path / "env-check-env-vars.yaml"
+    file_content = [
+        "dependencies: [numpy]",
+    ]
+    if env_vars:
+        variables_dict = {"MY_ENV_VAR": "My Value", "MY_OTHER_ENV_VAR": "Another Value"}
+        yaml_str = yaml.dump({"variables": variables_dict}, default_flow_style=False)
+        file_content.append(yaml_str)
+
+    with open(spec_file, "w") as f:
+        f.write("\n".join(file_content))
+
+    cmd = ["-n", env_name, "-f", spec_file, "--json"]
+    if no_env:
+        cmd += ["--no-env"]
+
+    res = helpers.create(*cmd)
+    assert res["success"]
+
+    packages = helpers.umamba_list("-n", env_name, "--json")
+    assert any(package["name"] == "numpy" for package in packages)
+
+    state_file_path = tmp_root_prefix / "envs" / env_name / "conda-meta" / "state"
+
+    if env_vars and not no_env:
+        assert state_file_path.exists()
+
+        with open(state_file_path) as f:
+            state_content = f.read()
+        assert (
+            '"env_vars":{"MY_ENV_VAR":"My Value","MY_OTHER_ENV_VAR":"Another Value"}'
+            in state_content
+        )
+    else:
+        assert not state_file_path.exists()
+
+
 @pytest.mark.parametrize("shared_pkgs_dirs", [True], indirect=True)
 @pytest.mark.parametrize("type", ("yaml", "classic", "explicit"))
 def test_multiple_spec_files(tmp_home, tmp_root_prefix, tmp_path, type):