Skip to content

Commit

Permalink
Environment map (#2967)
Browse files Browse the repository at this point in the history 
* Move env::copy yo utility and fix encoding

* Add set_env_map / update_env_map

* Fixup: `environ` is a define in msvc's stdlib.h

* Remove get_env_map use of to_upper

---------

Co-authored-by: Joël Lamotte (Klaim) <[email protected]>
  • Loading branch information
@AntoinePrv @Klaim
AntoinePrv and Klaim authored Nov 13, 2023
1 parent 0d42e81 commit b7b87f7
Show file tree
Hide file tree
Showing 7 changed files with 230 additions and 64 deletions.
4 changes: 2 additions & 2 deletions libmamba/include/mamba/core/activation.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,8 @@
#ifndef MAMBA_CORE_ACTIVATION_HPP
#define MAMBA_CORE_ACTIVATION_HPP

#include <map>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

Expand Down Expand Up @@ -101,7 +101,7 @@ namespace mamba
bool m_stack = false;
ActivationType m_action;

std::map<std::string, std::string> m_env;
std::unordered_map<std::string, std::string> m_env;
};

class PosixActivator : public Activator
Expand Down
2 changes: 0 additions & 2 deletions libmamba/include/mamba/core/environment.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@
#ifndef MAMBA_CORE_ENVIRONMENT_HPP
#define MAMBA_CORE_ENVIRONMENT_HPP

#include <map>
#include <string>
#include <vector>

Expand All @@ -20,7 +19,6 @@ namespace mamba::env

auto which(const std::string& exe, const std::string& override_path = "") -> fs::u8path;
auto which(const std::string& exe, const std::vector<fs::u8path>& search_paths) -> fs::u8path;
auto copy() -> std::map<std::string, std::string>;
auto platform() -> std::string;
auto home_directory() -> fs::u8path;
auto user_config_dir() -> fs::u8path;
Expand Down
38 changes: 37 additions & 1 deletion libmamba/include/mamba/util/environment.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,11 +9,47 @@

#include <optional>
#include <string>
#include <unordered_map>

namespace mamba::util
{
auto get_env(const std::string& key) -> std::optional<std::string>;
/**
* Get an environment variable encoded in UTF8.
*/
[[nodiscard]] auto get_env(const std::string& key) -> std::optional<std::string>;

/**
* Set an environment variable encoded in UTF8.
*/
void set_env(const std::string& key, const std::string& value);

/**
* Unset an environment variable encoded in UTF8.
*/
void unset_env(const std::string& key);

using environment_map = std::unordered_map<std::string, std::string>;

/**
* Return a map of all environment variables encoded in UTF8.
*
* This is useful if one in intersted to do an operatrion over all envrionment variables
* when their name is unknown.
*/
[[nodiscard]] auto get_env_map() -> environment_map;

/**
* Equivalent to calling set_env in a loop.
*
* This leaves environment variables not refered to in the map unmodified.
*/
void update_env_map(const environment_map& env);

/**
* Set the environment to be exactly the map given.
*
* This unset all environment variables not refered to in the map unmodified.
*/
void set_env_map(const environment_map& env);
}
#endif
2 changes: 1 addition & 1 deletion libmamba/src/core/activation.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ namespace mamba

Activator::Activator(const Context& context)
: m_context(context)
, m_env(env::copy())
, m_env(util::get_env_map())
{
}

Expand Down
42 changes: 0 additions & 42 deletions libmamba/src/core/environment.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -94,48 +94,6 @@ namespace mamba::env
return ""; // empty path
}

std::map<std::string, std::string> copy()
{
std::map<std::string, std::string> m;
#ifndef _WIN32
int i = 1;
const char* c = *environ;
for (; c; i++)
{
std::string_view s(c);
auto pos = s.find("=");
m[std::string(s.substr(0, pos))] = (pos != s.npos) ? std::string(s.substr(pos + 1)) : "";
c = *(environ + i);
}
#else

// inspired by
// https://github.com/gennaroprota/breath/blob/0709a9f0fe4e745b1d9fc44ab65d92853820b515
// /breath/environment/brt/dep/syst/windows/get_environment_map.cpp#L38-L80
char* start = GetEnvironmentStrings();
if (start == nullptr)
{
throw std::runtime_error("GetEnvironmentStrings() failed");
}

char* current = start;
while (*current != '\0')
{
std::string_view s = current;
auto pos = s.find("=");
assert(pos != std::string_view::npos);
std::string key = util::to_upper(s.substr(0, pos));
if (!key.empty())
{
std::string_view value = (pos != s.npos) ? s.substr(pos + 1) : "";
m[std::string(key)] = value;
}
current += s.size() + 1;
}
#endif
return m;
}

std::string platform()
{
#ifndef _WIN32
Expand Down
98 changes: 98 additions & 0 deletions libmamba/src/util/environment.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,11 +7,13 @@

#ifdef _WIN32

#include <cassert>
#include <mutex>
#include <stdexcept>

#include <fmt/format.h>
#include <Shlobj.h>
#include <Windows.h>

#include "mamba/util/environment.hpp"
#include "mamba/util/os_win.hpp"
Expand Down Expand Up @@ -92,6 +94,56 @@ namespace mamba::util
{
set_env(key, "");
}

namespace
{
struct Environ
{
Environ()
{
ptr = GetEnvironmentStringsW();
if (ptr == nullptr)
{
throw std::runtime_error("Fail to get environment");
}
}

~Environ()
{
FreeEnvironmentStringsW(ptr);
}

wchar_t* ptr = nullptr;
};
}

auto get_env_map() -> environment_map
{
static constexpr auto npos = std::wstring_view::npos;

auto env = environment_map();

auto raw_env = Environ();

wchar_t* current = raw_env.ptr;
while (*current != L'\0')
{
const auto expr = std::wstring_view(current);
const auto pos = expr.find(L'=');
assert(pos != npos);
std::string key = windows_encoding_to_utf8(expr.substr(0, pos));
if (!key.empty())
{
std::string value = windows_encoding_to_utf8(
(pos != npos) ? expr.substr(pos + 1) : L""
);
env.emplace(std::move(key), std::move(value));
}
current += expr.size() + 1;
}

return env;
}
}

#else // #ifdef _WIN32
Expand All @@ -103,6 +155,11 @@ namespace mamba::util

#include "mamba/util/environment.hpp"

extern "C"
{
extern char** environ; // Unix defined
}

namespace mamba::util
{
auto get_env(const std::string& key) -> std::optional<std::string>
Expand Down Expand Up @@ -133,6 +190,47 @@ namespace mamba::util
throw std::runtime_error(fmt::format(R"(Could not unset environment variable "{}")", key));
}
}

auto get_env_map() -> environment_map
{
// To keep the same signature between Unix and Windows (which is more convenient
// to work with), we have to copy strings.
// A more advanced slution could wrap a platform specific solution in a common return type.
auto env = environment_map();
for (std::size_t i = 0; environ[i]; ++i)
{
const auto expr = std::string_view(environ[i]);
const auto pos = expr.find('=');
env.emplace(
expr.substr(0, pos),
(pos != expr.npos) ? std::string(expr.substr(pos + 1)) : ""

);
}
return env;
}
}

#endif // #ifdef _WIN32

#include "mamba/util/environment.hpp"

namespace mamba::util
{
void update_env_map(const environment_map& env)
{
for (const auto& [name, val] : env)
{
set_env(name, val);
}
}

void set_env_map(const environment_map& env)
{
for (const auto& [name, val] : get_env_map())
{
unset_env(name);
}
update_env_map(env);
}
}
108 changes: 92 additions & 16 deletions libmamba/tests/src/util/test_environment.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,34 +8,110 @@

#include "mamba/util/environment.hpp"

using namespace mamba;
using namespace mamba::util;

TEST_SUITE("util::environment")
{
TEST_CASE("get_env")
{
CHECK_FALSE(util::get_env("VAR_THAT_DOES_NOT_EXIST_XYZ").has_value());
CHECK(util::get_env("PATH").has_value());
CHECK_FALSE(get_env("VAR_THAT_DOES_NOT_EXIST_XYZ").has_value());
CHECK(get_env("PATH").has_value());
}

TEST_CASE("set_env")
{
util::set_env("VAR_THAT_DOES_NOT_EXIST_XYZ", "VALUE");
CHECK_EQ(util::get_env("VAR_THAT_DOES_NOT_EXIST_XYZ"), "VALUE");
util::set_env(u8"VAR_私のにほんごわへたです", u8"😀");
CHECK_EQ(util::get_env(u8"VAR_私のにほんごわへたです"), u8"😀");
util::set_env(u8"VAR_私のにほんごわへたです", u8"hello");
CHECK_EQ(util::get_env(u8"VAR_私のにほんごわへたです"), u8"hello");
SUBCASE("ASCII")
{
const auto key = std::string(u8"VAR_THAT_DOES_NOT_EXIST_XYZ");
const auto value1 = std::string(u8"VALUE");
set_env(key, value1);
CHECK_EQ(get_env(key), value1);
const auto value2 = std::string(u8"VALUE_NEW");
set_env(key, value2);
CHECK_EQ(get_env(key), value2);
}

SUBCASE("UTF-8")
{
const auto key = std::string(u8"VAR_私のにほんごわへたです");
const auto value1 = std::string(u8"😀");
set_env(key, value1);
CHECK_EQ(get_env(key), value1);
const auto value2 = std::string(u8"🤗");
set_env(key, value2);
CHECK_EQ(get_env(key), value2);
}
}

TEST_CASE("unset_env")
{
CHECK_FALSE(util::get_env("VAR_THAT_DOES_NOT_EXIST_ZZZ").has_value());
util::unset_env("VAR_THAT_DOES_NOT_EXIST_ZZZ");
CHECK_FALSE(util::get_env("VAR_THAT_DOES_NOT_EXIST_ZZZ").has_value());
util::set_env("VAR_THAT_DOES_NOT_EXIST_ZZZ", "VALUE");
CHECK(util::get_env("VAR_THAT_DOES_NOT_EXIST_ZZZ").has_value());
util::unset_env("VAR_THAT_DOES_NOT_EXIST_ZZZ");
CHECK_FALSE(util::get_env("VAR_THAT_DOES_NOT_EXIST_ZZZ").has_value());
const auto key = std::string(u8"VAR_THAT_DOES_NOT_EXIST_ABC_😀");
CHECK_FALSE(get_env(key).has_value());
unset_env(key);
CHECK_FALSE(get_env(key).has_value());
set_env(key, "VALUE");
CHECK(get_env(key).has_value());
unset_env(key);
CHECK_FALSE(get_env(key).has_value());
}

TEST_CASE("get_env_map")
{
auto env = mamba::util::get_env_map();
CHECK_GT(env.size(), 0);
CHECK_EQ(env.count("VAR_THAT_MUST_NOT_EXIST_XYZ"), 0);
CHECK_EQ(env.count("PATH"), 1);

const auto key = std::string(u8"VAR_私のにほHelloわへたです");
const auto value = std::string(u8"😀");
set_env(key, value);
env = get_env_map();
CHECK_EQ(env.at(key), value);
}

TEST_CASE("update_env_map")
{
const auto key_inexistant = std::string(u8"CONDA😀");
const auto key_unchanged = std::string(u8"MAMBA😀");
const auto key_changed = std::string(u8"PIXI😀");

CHECK_FALSE(get_env(key_inexistant).has_value());
CHECK_FALSE(get_env(key_unchanged).has_value());
CHECK_FALSE(get_env(key_changed).has_value());

const auto val_set_1 = std::string(u8"a😀");
update_env_map({ { key_changed, val_set_1 }, { key_unchanged, val_set_1 } });
CHECK_EQ(get_env(key_inexistant), std::nullopt);
CHECK_EQ(get_env(key_unchanged), val_set_1);
CHECK_EQ(get_env(key_changed), val_set_1);

const auto val_set_2 = std::string(u8"b😀");
update_env_map({ { key_changed, val_set_2 } });
CHECK_EQ(get_env(key_inexistant), std::nullopt);
CHECK_EQ(get_env(key_unchanged), val_set_1);
CHECK_EQ(get_env(key_changed), val_set_2);
}

TEST_CASE("set_env_map")
{
const auto key_inexistant = std::string(u8"CONDA🤗");
const auto key_unchanged = std::string(u8"MAMBA🤗");
const auto key_changed = std::string(u8"PIXI🤗");

CHECK_FALSE(get_env(key_inexistant).has_value());
CHECK_FALSE(get_env(key_unchanged).has_value());
CHECK_FALSE(get_env(key_changed).has_value());

const auto val_set_1 = std::string(u8"a😀");
set_env_map({ { key_changed, val_set_1 }, { key_unchanged, val_set_1 } });
CHECK_EQ(get_env(key_inexistant), std::nullopt);
CHECK_EQ(get_env(key_unchanged), val_set_1);
CHECK_EQ(get_env(key_changed), val_set_1);

const auto val_set_2 = std::string(u8"b😀");
set_env_map({ { key_changed, val_set_2 } });
CHECK_EQ(get_env(key_inexistant), std::nullopt);
CHECK_EQ(get_env(key_unchanged), std::nullopt); // Difference with update_env_map
CHECK_EQ(get_env(key_changed), val_set_2);
}
}

0 comments on commit b7b87f7

Please sign in to comment.