Merge pull request #11494 from bryanhonof/bryanhonof.feat-mixenvironment-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"
-
+#include "environment-variables.hh"
-extern char * * environ __attribute__((weak));
 
 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",
+        std::string markdownError =
-                commandName, concatStringsSep("\n", subCommandTextLines));
+            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 =
+        evalState = std::allocate_shared<EvalState>(
-            std::allocate_shared<EvalState>(
+            traceable_allocator<EvalState>(), lookupPath, getEvalStore(), fetchSettings, evalSettings, getStore());
-                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({
+    addFlag(
-        .longName = "profile",
+        {.longName = "profile",
-        .description = "The profile to operate on.",
+         .description = "The profile to operate on.",
-        .labels = {"path"},
+         .labels = {"path"},
-        .handler = {&profile},
+         .handler = {&profile},
-        .completer = completePath
+         .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,
+    switchLink(profile2, createGeneration(*store, profile2, storePath));
-        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 {
+        std::visit(
-            [&](const BuiltPath::Opaque & bo) {
+            overloaded{
-                result.push_back(bo.path);
+                [&](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) {
+            buildable.raw());
-                for (auto & output : bfd.outputs) {
-                    result.push_back(output.second);
-                }
-            },
-        }, 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() {
+void MixEnvironment::setEnviron()
-    if (ignoreEnvironment) {
+{
-        if (!unset.empty())
+    if (ignoreEnvironment && !unsetVars.empty())
-            throw UsageError("--unset does not make sense with --ignore-environment");
+        throw UsageError("--unset-env-var does not make sense with --ignore-env");
 
-        for (const auto & var : keep) {
+    if (!ignoreEnvironment && !keepVars.empty())
-            auto val = getenv(var.c_str());
+        throw UsageError("--keep-env-var does not make sense without --ignore-env");
-            if (val) stringsEnv.emplace_back(fmt("%s=%s", var.c_str(), val));
-        }
 
-        vectorEnv = stringsToCharPtrs(stringsEnv);
+    auto env = getEnv();
-        environ = vectorEnv.data();
-    } else {
-        if (!keep.empty())
-            throw UsageError("--keep does not make sense without --ignore-environment");
 
-        for (const auto & var : unset)
+    if (ignoreEnvironment)
-            unsetenv(var.c_str());
+        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(
+    Installables parseInstallables(ref<Store> store, std::vector<std::string> ss);
-        ref<Store> store, std::vector<std::string> ss);
 
-    ref<Installable> parseInstallable(
+    ref<Installable> parseInstallable(ref<Store> store, const std::string & installable);
-        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,
+    RegisterCommand(std::vector<std::string> && name, std::function<ref<Command>()> command)
-        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;
+    StringSet keepVars;
-    Strings stringsEnv;
+    StringSet unsetVars;
-    std::vector<char*> vectorEnv;
+    std::map<std::string, std::string> setVars;
     bool ignoreEnvironment;
 
     MixEnvironment();
 
     /***
-     * Modify global environ based on `ignoreEnvironment`, `keep`, and
+     * Modify global environ based on `ignoreEnvironment`, `keep`,
-     * `unset`. It's expected that exec will be called before this class
+     * `unset`, and `added`. It's expected that exec will be called
-     * goes out of scope, otherwise `environ` will become invalid.
+     * 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,
+    ref<Store> store, const StorePath & beforePath, const StorePath & afterPath, std::string_view indent);
-    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.
  */

tests/functional/flakes/develop.sh

@@ -5,11 +5,11 @@ source ../common.sh
 TODO_NixOS
 
 clearStore
-rm -rf $TEST_HOME/.cache $TEST_HOME/.config $TEST_HOME/.local
+rm -rf "$TEST_HOME/.cache" "$TEST_HOME/.config" "$TEST_HOME/.local"
 
 # Create flake under test.
-cp ../shell-hello.nix "${config_nix}" $TEST_HOME/
+cp ../shell-hello.nix "$config_nix" "$TEST_HOME/"
-cat <<EOF >$TEST_HOME/flake.nix
+cat <<EOF >"$TEST_HOME/flake.nix"
 {
     inputs.nixpkgs.url = "$TEST_HOME/nixpkgs";
     outputs = {self, nixpkgs}: {
@@ -24,13 +24,13 @@ cat <<EOF >$TEST_HOME/flake.nix
 EOF
 
 # Create fake nixpkgs flake.
-mkdir -p $TEST_HOME/nixpkgs
+mkdir -p "$TEST_HOME/nixpkgs"
-cp "${config_nix}" ../shell.nix $TEST_HOME/nixpkgs
+cp "${config_nix}" ../shell.nix "$TEST_HOME/nixpkgs"
 
 # `config.nix` cannot be gotten via build dir / env var (runs afoul pure eval). Instead get from flake.
 removeBuildDirRef "$TEST_HOME/nixpkgs"/*.nix
 
-cat <<EOF >$TEST_HOME/nixpkgs/flake.nix
+cat <<EOF >"$TEST_HOME/nixpkgs/flake.nix"
 {
     outputs = {self}: {
       legacyPackages.$system.bashInteractive = (import ./shell.nix {}).bashInteractive;
@@ -38,7 +38,7 @@ cat <<EOF >$TEST_HOME/nixpkgs/flake.nix
 }
 EOF
 
-cd $TEST_HOME
+cd "$TEST_HOME"
 
 # Test whether `nix develop` passes through environment variables.
 [[ "$(
@@ -47,13 +47,86 @@ echo "\$ENVVAR"
 EOF
 )" = "a" ]]
 
-# Test whether `nix develop --ignore-environment` does _not_ pass through environment variables.
+# Test whether `nix develop --ignore-env` does _not_ pass through environment variables.
 [[ -z "$(
-    ENVVAR=a nix develop --ignore-environment --no-write-lock-file .#hello <<EOF
+    ENVVAR=a nix develop --ignore-env --no-write-lock-file .#hello <<EOF
 echo "\$ENVVAR"
 EOF
 )" ]]
 
+# Test wether `--keep-env-var` keeps the environment variable.
+(
+  expect='BAR'
+  got="$(FOO='BAR' nix develop --ignore-env --keep-env-var FOO --no-write-lock-file .#hello <<EOF
+echo "\$FOO"
+EOF
+)"
+  [[ "$got" == "$expect" ]]
+)
+
+# Test wether duplicate `--keep-env-var` keeps the environment variable.
+(
+  expect='BAR'
+  got="$(FOO='BAR' nix develop --ignore-env --keep-env-var FOO --keep-env-var FOO --no-write-lock-file .#hello <<EOF
+echo "\$FOO"
+EOF
+)"
+  [[ "$got" == "$expect" ]]
+)
+
+# Test wether `--set-env-var` sets the environment variable.
+(
+  expect='BAR'
+  got="$(nix develop --ignore-env --set-env-var FOO 'BAR' --no-write-lock-file .#hello <<EOF
+echo "\$FOO"
+EOF
+)"
+  [[ "$got" == "$expect" ]]
+)
+
+# Test that `--set-env-var` overwrites previously set variables.
+(
+  expect='BLA'
+  got="$(FOO='BAR' nix develop --set-env-var FOO 'BLA' --no-write-lock-file .#hello <<EOF
+echo "\$FOO"
+EOF
+)"
+  [[ "$got" == "$expect" ]]
+)
+
+# Test that multiple `--set-env-var` work.
+(
+  expect='BARFOO'
+  got="$(nix develop --set-env-var FOO 'BAR' --set-env-var BAR 'FOO' --no-write-lock-file .#hello <<EOF | tr -d '\n'
+echo "\$FOO"
+echo "\$BAR"
+EOF
+)"
+  [[ "$got" == "$expect" ]]
+)
+
+# Check that we throw an error when `--keep-env-var` is used without `--ignore-env`.
+expectStderr 1 nix develop --keep-env-var FOO .#hello |
+  grepQuiet "error: --keep-env-var does not make sense without --ignore-env"
+
+# Check that we throw an error when `--unset-env-var` is used with `--ignore-env`.
+expectStderr 1 nix develop --ignore-env --unset-env-var FOO .#hello |
+  grepQuiet "error: --unset-env-var does not make sense with --ignore-env"
+
+# Test wether multiple occurances of `--set-env-var` throws.
+expectStderr 1 nix develop --set-env-var FOO 'BAR' --set-env-var FOO 'BLA' --no-write-lock-file .#hello |
+  grepQuiet "error: Duplicate definition of environment variable 'FOO' with '--set-env-var' is ambiguous"
+
+# Test wether similar `--unset-env-var` and `--set-env-var` throws.
+expectStderr 1 nix develop --set-env-var FOO 'BAR' --unset-env-var FOO --no-write-lock-file .#hello |
+  grepQuiet "error: Cannot unset environment variable 'FOO' that is set with '--set-env-var'"
+
+expectStderr 1 nix develop --unset-env-var FOO --set-env-var FOO 'BAR' --no-write-lock-file .#hello |
+  grepQuiet "error: Cannot set environment variable 'FOO' that is unset with '--unset-env-var'"
+
+# Check that multiple `--ignore-env`'s are okay.
+expectStderr 0 nix develop --ignore-env --set-env-var FOO 'BAR' --ignore-env .#hello
+
 # Determine the bashInteractive executable.
 nix build --no-write-lock-file './nixpkgs#bashInteractive' --out-link ./bash-interactive
 BASH_INTERACTIVE_EXECUTABLE="$PWD/bash-interactive/bin/bash"
@@ -67,7 +140,7 @@ EOF
 
 # Test whether `nix develop` with ignore environment sets `SHELL` to nixpkgs#bashInteractive shell.
 [[ "$(
-    SHELL=custom nix develop --ignore-environment --no-write-lock-file .#hello <<EOF
+    SHELL=custom nix develop --ignore-env --no-write-lock-file .#hello <<EOF
 echo "\$SHELL"
 EOF
 )" -ef "$BASH_INTERACTIVE_EXECUTABLE" ]]