diff --git a/include/gz/utils/Environment.hh b/include/gz/utils/Environment.hh index 9a7a062..c282920 100644 --- a/include/gz/utils/Environment.hh +++ b/include/gz/utils/Environment.hh @@ -23,6 +23,7 @@ #include #include +#include namespace gz { @@ -78,6 +79,22 @@ bool GZ_UTILS_VISIBLE clearenv(); /// \brief Type alias for a collection of environment variables using EnvironmentMap = std::unordered_map; +/// \brief Type alias for a collection of environment variables +/// Each entry is of the form KEY=VAL +using EnvironmentStrings = std::vector; + +/// \brief Convert a vector of environment variables to a map +/// +/// \param[in] _envStrings Vector collection of environment variables. +/// \return Mapped collection of environment variables. +EnvironmentMap GZ_UTILS_VISIBLE envStringsToMap(const EnvironmentStrings &_envStrings); + +/// \brief Convert a map of environment variables to a vector +/// +/// \param[in] _envMap Collection of mapped environment variables. +/// \return Vector collection of environment variables. +EnvironmentStrings GZ_UTILS_VISIBLE envMapToStrings(const EnvironmentMap &_envMap); + /// \brief Retrieve all current environment variables /// /// Note: This function is not thread-safe and should not be called @@ -108,6 +125,7 @@ bool GZ_UTILS_VISIBLE setenv(const EnvironmentMap &_vars); /// \return A string containing all environment variables /// NOLINTNEXTLINE - This is incorrectly parsed as a global variable std::string GZ_UTILS_VISIBLE printenv(); + } // namespace GZ_UTILS_VERSION_NAMESPACE } // namespace utils } // namespace gz diff --git a/include/gz/utils/Subprocess.hh b/include/gz/utils/Subprocess.hh index dee00fd..843879d 100644 --- a/include/gz/utils/Subprocess.hh +++ b/include/gz/utils/Subprocess.hh @@ -72,6 +72,22 @@ class Subprocess this->Create(); } + /// \brief Constructor + /// + /// This variant will spawn a subprocess that uses the user-specified + /// environment + /// + /// \param[in] _commandLine set of arguments starting with an executable + /// used to spawn the subprocess + /// \param[in] _environment environment variables to set in the spawned + /// subprocess + public: Subprocess(const std::vector &_commandLine, + const gz::utils::EnvironmentStrings &_environment): + Subprocess(_commandLine, gz::utils::envStringsToMap(_environment)) + { + } + + private: void Create() { if (this->process != nullptr) diff --git a/src/Environment.cc b/src/Environment.cc index 4d86e0b..a5b47e8 100644 --- a/src/Environment.cc +++ b/src/Environment.cc @@ -136,14 +136,6 @@ EnvironmentMap env() { EnvironmentMap ret; - // Helper function to split KEY=VAL - auto split = [](const std::string &_inp) - { - return std::make_pair( - _inp.substr(0, _inp.find('=')), - _inp.substr(_inp.find('=') + 1)); - }; - char **currentEnv = nullptr; #ifdef _WIN32 currentEnv = *__p__environ(); @@ -155,11 +147,12 @@ EnvironmentMap env() if (currentEnv == nullptr) return {}; + std::vector envStrings; for (; *currentEnv; ++currentEnv) { - ret.emplace(split(*currentEnv)); + envStrings.emplace_back(*currentEnv); } - return ret; + return envStringsToMap(envStrings); } ///////////////////////////////////////////////// @@ -174,20 +167,41 @@ bool setenv(const EnvironmentMap &_vars) } ///////////////////////////////////////////////// -std::string printenv() +EnvironmentMap envStringsToMap(const EnvironmentStrings &_envStrings) { - std::string ret; - // Variables are in an unordered_map as we generally don't - // care, but for printing sort for consistent display - auto currentEnv = env(); + EnvironmentMap ret; + for (const auto &pair : _envStrings) + { + auto eqPos = pair.find('='); + if (eqPos != std::string::npos) + { + ret.emplace(pair.substr(0, eqPos), pair.substr(eqPos + 1)); + } + } + return ret; +} + +///////////////////////////////////////////////// +EnvironmentStrings envMapToStrings(const EnvironmentMap &_envMap) +{ + EnvironmentStrings ret; auto sorted = std::vector>( - currentEnv.begin(), currentEnv.end()); + _envMap.begin(), _envMap.end()); std::sort(sorted.begin(), sorted.end()); - for (const auto &[key, value] : sorted) + for (auto [key, value] : sorted) + { + ret.push_back(key + "=" + value); + } + return ret; +} + +///////////////////////////////////////////////// +std::string printenv() +{ + std::string ret; + for (const auto &entry : envMapToStrings(env())) { - ret.append(key); - ret.append("="); - ret.append(value); + ret.append(entry); ret.append("\n"); } return ret; diff --git a/src/Environment_TEST.cc b/src/Environment_TEST.cc index 6d09923..496dddb 100644 --- a/src/Environment_TEST.cc +++ b/src/Environment_TEST.cc @@ -175,3 +175,39 @@ TEST(Environment, printenv) EXPECT_EQ(gz::utils::printenv(), "GZ_BAR_KEY=GZ_BAR_VAL\nGZ_BAZ_KEY=GZ_BAZ_VAL\nGZ_FOO_KEY=GZ_FOO_VAL\n"); } + +///////////////////////////////////////////////// +TEST(Environment, envStringsToMap) +{ + gz::utils::EnvironmentStrings strings; + strings.emplace_back("GZ_FOO_KEY=GZ_FOO_VAL"); + strings.emplace_back("GZ_BAR_KEY=GZ_BAR_VAL"); + strings.emplace_back("GZ_BAZ_KEY=GZ_BAZ_VAL"); + strings.emplace_back("BAD_KEY"); + + { + auto envMap = gz::utils::envStringsToMap(strings); + EXPECT_EQ(3u, envMap.size()); + EXPECT_EQ("GZ_FOO_VAL", envMap["GZ_FOO_KEY"]); + EXPECT_EQ("GZ_BAR_VAL", envMap["GZ_BAR_KEY"]); + EXPECT_EQ("GZ_BAZ_VAL", envMap["GZ_BAZ_KEY"]); + } +} + +///////////////////////////////////////////////// +TEST(Environment, envMapToStrings) +{ + gz::utils::EnvironmentMap envMap; + envMap.insert({{"GZ_FOO_KEY"}, {"GZ_FOO_VAL"}}); + envMap.insert({{"GZ_BAR_KEY"}, {"GZ_BAR_VAL"}}); + envMap.insert({{"GZ_BAZ_KEY"}, {"GZ_BAZ_VAL"}}); + + { + auto envStrings = gz::utils::envMapToStrings(envMap); + + EXPECT_EQ(3u, envStrings.size()); + EXPECT_EQ("GZ_BAR_KEY=GZ_BAR_VAL", envStrings[0]); + EXPECT_EQ("GZ_BAZ_KEY=GZ_BAZ_VAL", envStrings[1]); + EXPECT_EQ("GZ_FOO_KEY=GZ_FOO_VAL", envStrings[2]); + } +} diff --git a/test/integration/subprocess_TEST.cc b/test/integration/subprocess_TEST.cc index 911ac60..47c8976 100644 --- a/test/integration/subprocess_TEST.cc +++ b/test/integration/subprocess_TEST.cc @@ -138,7 +138,23 @@ TEST(Subprocess, Environment) { // Passing an empty map as the second arg clears the environment auto proc = Subprocess( - {kExecutablePath, "--output=cout", "--environment"}, {}); + {kExecutablePath, "--output=cout", "--environment"}, + gz::utils::EnvironmentMap()); + // Block until the executable is done + auto ret = proc.Join(); + EXPECT_EQ(0u, ret); + + auto cout = proc.Stdout(); + EXPECT_EQ(std::string::npos, cout.find("GZ_FOO_KEY=GZ_FOO_VAL")); + EXPECT_EQ(std::string::npos, cout.find("GZ_BAR_KEY=GZ_BAR_VAL")); + EXPECT_EQ(std::string::npos, cout.find("GZ_BAZ_KEY=GZ_BAZ_VAL")); + } + + { + // Passing an empty vector as the second arg clears the environment + auto proc = Subprocess( + {kExecutablePath, "--output=cout", "--environment"}, + gz::utils::EnvironmentStrings()); // Block until the executable is done auto ret = proc.Join(); EXPECT_EQ(0u, ret); @@ -152,9 +168,10 @@ TEST(Subprocess, Environment) { // Passing a map sets those variables, clearing the rest auto proc = Subprocess( - {kExecutablePath, "--output=cout", "--environment"}, { - {"QUX_KEY", "QUX_VAL"} - }); + {kExecutablePath, "--output=cout", "--environment"}, + gz::utils::EnvironmentMap{ + {"QUX_KEY", "QUX_VAL"} + }); // Block until the executable is done auto ret = proc.Join(); EXPECT_EQ(0u, ret); @@ -165,4 +182,23 @@ TEST(Subprocess, Environment) EXPECT_EQ(std::string::npos, cout.find("GZ_BAZ_KEY=GZ_BAZ_VAL")); EXPECT_NE(std::string::npos, cout.find("QUX_KEY=QUX_VAL")); } + + { + // Passing a map sets those variables, clearing the rest + auto proc = Subprocess( + {kExecutablePath, "--output=cout", "--environment"}, + gz::utils::EnvironmentStrings{ + {"QUX_KEY=QUX_VAL"} + }); + // Block until the executable is done + auto ret = proc.Join(); + EXPECT_EQ(0u, ret); + + auto cout = proc.Stdout(); + EXPECT_EQ(std::string::npos, cout.find("GZ_FOO_KEY=GZ_FOO_VAL")); + EXPECT_EQ(std::string::npos, cout.find("GZ_BAR_KEY=GZ_BAR_VAL")); + EXPECT_EQ(std::string::npos, cout.find("GZ_BAZ_KEY=GZ_BAZ_VAL")); + EXPECT_NE(std::string::npos, cout.find("QUX_KEY=QUX_VAL")); + } + }