Skip to content

Commit

Permalink
Add helper functions and delegating constructor
Browse files Browse the repository at this point in the history
Signed-off-by: Michael Carroll <mjcarroll@intrinsic.ai>
  • Loading branch information
mjcarroll committed Nov 8, 2023
1 parent 8b36418 commit 4280e46
Show file tree
Hide file tree
Showing 5 changed files with 144 additions and 24 deletions.
18 changes: 18 additions & 0 deletions include/gz/utils/Environment.hh
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@

#include <string>
#include <unordered_map>
#include <vector>

namespace gz
{
Expand Down Expand Up @@ -78,6 +79,22 @@ bool GZ_UTILS_VISIBLE clearenv();
/// \brief Type alias for a collection of environment variables
using EnvironmentMap = std::unordered_map<std::string, std::string>;

/// \brief Type alias for a collection of environment variables
/// Each entry is of the form KEY=VAL
using EnvironmentStrings = std::vector<std::string>;

/// \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
Expand Down Expand Up @@ -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
Expand Down
16 changes: 16 additions & 0 deletions include/gz/utils/Subprocess.hh
Original file line number Diff line number Diff line change
Expand Up @@ -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<std::string> &_commandLine,
const gz::utils::EnvironmentStrings &_environment):
Subprocess(_commandLine, gz::utils::envStringsToMap(_environment))
{
}


private: void Create()
{
if (this->process != nullptr)
Expand Down
54 changes: 34 additions & 20 deletions src/Environment.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand All @@ -155,11 +147,12 @@ EnvironmentMap env()
if (currentEnv == nullptr)
return {};

std::vector<std::string> envStrings;
for (; *currentEnv; ++currentEnv)
{
ret.emplace(split(*currentEnv));
envStrings.emplace_back(*currentEnv);
}
return ret;
return envStringsToMap(envStrings);
}

/////////////////////////////////////////////////
Expand All @@ -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<std::pair<std::string, std::string>>(
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;
Expand Down
36 changes: 36 additions & 0 deletions src/Environment_TEST.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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]);
}
}
44 changes: 40 additions & 4 deletions test/integration/subprocess_TEST.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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);
Expand All @@ -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"));
}

}

0 comments on commit 4280e46

Please sign in to comment.