Merge pull request #11494 from bryanhonof/bryanhonof.feat-mixenvironm…

…ent-set-env

Add flag `set-env-var` to `MixEnvironment`, and rename old ones

maintainers/flake-module.nix

@@ -28,8 +28,6 @@
             ''^src/build-remote/build-remote\.cc$''
             ''^src/libcmd/built-path\.cc$''
             ''^src/libcmd/built-path\.hh$''
-            ''^src/libcmd/command\.cc$''
-            ''^src/libcmd/command\.hh$''
             ''^src/libcmd/common-eval-args\.cc$''
             ''^src/libcmd/common-eval-args\.hh$''
             ''^src/libcmd/editor-for\.cc$''
@@ -547,7 +545,6 @@
             ''^tests/functional/flakes/absolute-paths\.sh$''
             ''^tests/functional/flakes/check\.sh$''
             ''^tests/functional/flakes/config\.sh$''
-            ''^tests/functional/flakes/develop\.sh$''
             ''^tests/functional/flakes/flakes\.sh$''
             ''^tests/functional/flakes/follow-paths\.sh$''
             ''^tests/functional/flakes/prefetch\.sh$''

src/libcmd/command.cc

@@ -1,3 +1,4 @@
+#include <algorithm>
 #include <nlohmann/json.hpp>
 
 #include "command.hh"
@@ -9,8 +10,7 @@
 #include "profiles.hh"
 #include "repl.hh"
 #include "strings.hh"
-
-extern char * * environ __attribute__((weak));
+#include "environment-variables.hh"
 
 namespace nix {
 
@@ -23,7 +23,8 @@ nix::Commands RegisterCommand::getCommandsFor(const std::vector<std::string> & p
         if (name.size() == prefix.size() + 1) {
             bool equal = true;
             for (size_t i = 0; i < prefix.size(); ++i)
-                if (name[i] != prefix[i]) equal = false;
+                if (name[i] != prefix[i])
+                    equal = false;
             if (equal)
                 res.insert_or_assign(name[prefix.size()], command);
         }
@@ -42,16 +43,16 @@ void NixMultiCommand::run()
         std::set<std::string> subCommandTextLines;
         for (auto & [name, _] : commands)
             subCommandTextLines.insert(fmt("- `%s`", name));
-        std::string markdownError = fmt("`nix %s` requires a sub-command. Available sub-commands:\n\n%s\n",
-                commandName, concatStringsSep("\n", subCommandTextLines));
+        std::string markdownError =
+            fmt("`nix %s` requires a sub-command. Available sub-commands:\n\n%s\n",
+                commandName,
+                concatStringsSep("\n", subCommandTextLines));
         throw UsageError(renderMarkdownToTerminal(markdownError));
     }
     command->second->run();
 }
 
-StoreCommand::StoreCommand()
-{
-}
+StoreCommand::StoreCommand() {}
 
 ref<Store> StoreCommand::getStore()
 {
@@ -126,10 +127,8 @@ ref<Store> EvalCommand::getEvalStore()
 ref<EvalState> EvalCommand::getEvalState()
 {
     if (!evalState) {
-        evalState =
-            std::allocate_shared<EvalState>(
-                traceable_allocator<EvalState>(),
-                lookupPath, getEvalStore(), fetchSettings, evalSettings, getStore());
+        evalState = std::allocate_shared<EvalState>(
+            traceable_allocator<EvalState>(), lookupPath, getEvalStore(), fetchSettings, evalSettings, getStore());
 
         evalState->repair = repair;
 
@@ -144,7 +143,8 @@ MixOperateOnOptions::MixOperateOnOptions()
 {
     addFlag({
         .longName = "derivation",
-        .description = "Operate on the [store derivation](@docroot@/glossary.md#gloss-store-derivation) rather than its outputs.",
+        .description =
+            "Operate on the [store derivation](@docroot@/glossary.md#gloss-store-derivation) rather than its outputs.",
         .category = installablesCategory,
         .handler = {&operateOn, OperateOn::Derivation},
     });
@@ -233,46 +233,48 @@ void StorePathCommand::run(ref<Store> store, StorePaths && storePaths)
 
 MixProfile::MixProfile()
 {
-    addFlag({
-        .longName = "profile",
-        .description = "The profile to operate on.",
-        .labels = {"path"},
-        .handler = {&profile},
-        .completer = completePath
-    });
+    addFlag(
+        {.longName = "profile",
+         .description = "The profile to operate on.",
+         .labels = {"path"},
+         .handler = {&profile},
+         .completer = completePath});
 }
 
 void MixProfile::updateProfile(const StorePath & storePath)
 {
-    if (!profile) return;
+    if (!profile)
+        return;
     auto store = getStore().dynamic_pointer_cast<LocalFSStore>();
-    if (!store) throw Error("'--profile' is not supported for this Nix store");
+    if (!store)
+        throw Error("'--profile' is not supported for this Nix store");
     auto profile2 = absPath(*profile);
-    switchLink(profile2,
-        createGeneration(*store, profile2, storePath));
+    switchLink(profile2, createGeneration(*store, profile2, storePath));
 }
 
 void MixProfile::updateProfile(const BuiltPaths & buildables)
 {
-    if (!profile) return;
+    if (!profile)
+        return;
 
     StorePaths result;
 
     for (auto & buildable : buildables) {
-        std::visit(overloaded {
-            [&](const BuiltPath::Opaque & bo) {
-                result.push_back(bo.path);
+        std::visit(
+            overloaded{
+                [&](const BuiltPath::Opaque & bo) { result.push_back(bo.path); },
+                [&](const BuiltPath::Built & bfd) {
+                    for (auto & output : bfd.outputs) {
+                        result.push_back(output.second);
+                    }
+                },
             },
-            [&](const BuiltPath::Built & bfd) {
-                for (auto & output : bfd.outputs) {
-                    result.push_back(output.second);
-                }
-            },
-        }, buildable.raw());
+            buildable.raw());
     }
 
     if (result.size() != 1)
-        throw UsageError("'--profile' requires that the arguments produce a single store path, but there are %d", result.size());
+        throw UsageError(
+            "'--profile' requires that the arguments produce a single store path, but there are %d", result.size());
 
     updateProfile(result[0]);
 }
@@ -282,51 +284,85 @@ MixDefaultProfile::MixDefaultProfile()
     profile = getDefaultProfile();
 }
 
-MixEnvironment::MixEnvironment() : ignoreEnvironment(false)
+MixEnvironment::MixEnvironment()
+    : ignoreEnvironment(false)
 {
     addFlag({
-        .longName = "ignore-environment",
+        .longName = "ignore-env",
+        .aliases = {"ignore-environment"},
         .shortName = 'i',
-        .description = "Clear the entire environment (except those specified with `--keep`).",
+        .description = "Clear the entire environment, except for those specified with `--keep-env-var`.",
+        .category = environmentVariablesCategory,
         .handler = {&ignoreEnvironment, true},
     });
 
     addFlag({
-        .longName = "keep",
+        .longName = "keep-env-var",
+        .aliases = {"keep"},
         .shortName = 'k',
-        .description = "Keep the environment variable *name*.",
+        .description = "Keep the environment variable *name*, when using `--ignore-env`.",
+        .category = environmentVariablesCategory,
         .labels = {"name"},
-        .handler = {[&](std::string s) { keep.insert(s); }},
+        .handler = {[&](std::string s) { keepVars.insert(s); }},
     });
 
     addFlag({
-        .longName = "unset",
+        .longName = "unset-env-var",
+        .aliases = {"unset"},
         .shortName = 'u',
         .description = "Unset the environment variable *name*.",
+        .category = environmentVariablesCategory,
         .labels = {"name"},
-        .handler = {[&](std::string s) { unset.insert(s); }},
+        .handler = {[&](std::string name) {
+            if (setVars.contains(name))
+                throw UsageError("Cannot unset environment variable '%s' that is set with '%s'", name, "--set-env-var");
+
+            unsetVars.insert(name);
+        }},
+    });
+
+    addFlag({
+        .longName = "set-env-var",
+        .shortName = 's',
+        .description = "Sets an environment variable *name* with *value*.",
+        .category = environmentVariablesCategory,
+        .labels = {"name", "value"},
+        .handler = {[&](std::string name, std::string value) {
+            if (unsetVars.contains(name))
+                throw UsageError(
+                    "Cannot set environment variable '%s' that is unset with '%s'", name, "--unset-env-var");
+
+            if (setVars.contains(name))
+                throw UsageError(
+                    "Duplicate definition of environment variable '%s' with '%s' is ambiguous", name, "--set-env-var");
+
+            setVars.insert_or_assign(name, value);
+        }},
     });
 }
 
-void MixEnvironment::setEnviron() {
-    if (ignoreEnvironment) {
-        if (!unset.empty())
-            throw UsageError("--unset does not make sense with --ignore-environment");
+void MixEnvironment::setEnviron()
+{
+    if (ignoreEnvironment && !unsetVars.empty())
+        throw UsageError("--unset-env-var does not make sense with --ignore-env");
 
-        for (const auto & var : keep) {
-            auto val = getenv(var.c_str());
-            if (val) stringsEnv.emplace_back(fmt("%s=%s", var.c_str(), val));
-        }
+    if (!ignoreEnvironment && !keepVars.empty())
+        throw UsageError("--keep-env-var does not make sense without --ignore-env");
 
-        vectorEnv = stringsToCharPtrs(stringsEnv);
-        environ = vectorEnv.data();
-    } else {
-        if (!keep.empty())
-            throw UsageError("--keep does not make sense without --ignore-environment");
+    auto env = getEnv();
 
-        for (const auto & var : unset)
-            unsetenv(var.c_str());
-    }
+    if (ignoreEnvironment)
+        std::erase_if(env, [&](const auto & var) { return !keepVars.contains(var.first); });
+
+    for (const auto & [name, value] : setVars)
+        env[name] = value;
+
+    if (!unsetVars.empty())
+        std::erase_if(env, [&](const auto & var) { return unsetVars.contains(var.first); });
+
+    replaceEnv(env);
+
+    return;
 }
 
 }

src/libcmd/command.hh

@@ -13,7 +13,7 @@ namespace nix {
 
 extern std::string programPath;
 
-extern char * * savedArgv;
+extern char ** savedArgv;
 
 class EvalState;
 struct Pos;
@@ -24,7 +24,8 @@ static constexpr Command::Category catSecondary = 100;
 static constexpr Command::Category catUtility = 101;
 static constexpr Command::Category catNixInstallation = 102;
 
-static constexpr auto installablesCategory = "Options that change the interpretation of [installables](@docroot@/command-ref/new-cli/nix.md#installables)";
+static constexpr auto installablesCategory =
+    "Options that change the interpretation of [installables](@docroot@/command-ref/new-cli/nix.md#installables)";
 
 struct NixMultiCommand : MultiCommand, virtual Command
 {
@@ -112,7 +113,9 @@ struct MixFlakeOptions : virtual Args, EvalCommand
      * arguments) so that the completions for these flags can use them.
      */
     virtual std::vector<FlakeRef> getFlakeRefsForCompletion()
-    { return {}; }
+    {
+        return {};
+    }
 };
 
 struct SourceExprCommand : virtual Args, MixFlakeOptions
@@ -122,11 +125,9 @@ struct SourceExprCommand : virtual Args, MixFlakeOptions
 
     SourceExprCommand();
 
-    Installables parseInstallables(
-        ref<Store> store, std::vector<std::string> ss);
+    Installables parseInstallables(ref<Store> store, std::vector<std::string> ss);
 
-    ref<Installable> parseInstallable(
-        ref<Store> store, const std::string & installable);
+    ref<Installable> parseInstallable(ref<Store> store, const std::string & installable);
 
     virtual Strings getDefaultFlakeAttrPaths();
 
@@ -272,10 +273,10 @@ struct RegisterCommand
     typedef std::map<std::vector<std::string>, std::function<ref<Command>()>> Commands;
     static Commands * commands;
 
-    RegisterCommand(std::vector<std::string> && name,
-        std::function<ref<Command>()> command)
+    RegisterCommand(std::vector<std::string> && name, std::function<ref<Command>()> command)
     {
-        if (!commands) commands = new Commands;
+        if (!commands)
+            commands = new Commands;
         commands->emplace(name, command);
     }
 
@@ -285,13 +286,13 @@ struct RegisterCommand
 template<class T>
 static RegisterCommand registerCommand(const std::string & name)
 {
-    return RegisterCommand({name}, [](){ return make_ref<T>(); });
+    return RegisterCommand({name}, []() { return make_ref<T>(); });
 }
 
 template<class T>
 static RegisterCommand registerCommand2(std::vector<std::string> && name)
 {
-    return RegisterCommand(std::move(name), [](){ return make_ref<T>(); });
+    return RegisterCommand(std::move(name), []() { return make_ref<T>(); });
 }
 
 struct MixProfile : virtual StoreCommand
@@ -313,19 +314,21 @@ struct MixDefaultProfile : MixProfile
     MixDefaultProfile();
 };
 
-struct MixEnvironment : virtual Args {
+struct MixEnvironment : virtual Args
+{
 
-    StringSet keep, unset;
-    Strings stringsEnv;
-    std::vector<char*> vectorEnv;
+    StringSet keepVars;
+    StringSet unsetVars;
+    std::map<std::string, std::string> setVars;
     bool ignoreEnvironment;
 
     MixEnvironment();
 
     /***
-     * Modify global environ based on `ignoreEnvironment`, `keep`, and
-     * `unset`. It's expected that exec will be called before this class
-     * goes out of scope, otherwise `environ` will become invalid.
+     * Modify global environ based on `ignoreEnvironment`, `keep`,
+     * `unset`, and `added`. It's expected that exec will be called
+     * before this class goes out of scope, otherwise `environ` will
+     * become invalid.
      */
     void setEnviron();
 };
@@ -349,9 +352,6 @@ void completeFlakeRefWithFragment(
 std::string showVersions(const std::set<std::string> & versions);
 
 void printClosureDiff(
-    ref<Store> store,
-    const StorePath & beforePath,
-    const StorePath & afterPath,
-    std::string_view indent);
+    ref<Store> store, const StorePath & beforePath, const StorePath & afterPath, std::string_view indent);
 
 }

src/libutil/environment-variables.hh

@@ -13,6 +13,8 @@
 
 namespace nix {
 
+static constexpr auto environmentVariablesCategory = "Options that change environment variables";
+
 /**
  * @return an environment variable.
  */