unix/modos: Implement putenv and unsetenv to complement getenv.

CPython also has os.environ, which should be used instead of os.getenv() due to caching in the os.environ mapping. But for MicroPython it makes sense to only implement the basic underlying methods, ie getenv/putenv/ unsetenv.
micropython · Feb 4, 2020 · f1b6e6b · f1b6e6b
1 parent 83439e3
commit f1b6e6b
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 0 deletions.
diff --git a/ports/unix/modos.c b/ports/unix/modos.c
@@ -180,6 +180,41 @@ STATIC mp_obj_t mod_os_getenv(mp_obj_t var_in) {
 }
 MP_DEFINE_CONST_FUN_OBJ_1(mod_os_getenv_obj, mod_os_getenv);
 
+STATIC mp_obj_t mod_os_putenv(mp_obj_t key_in, mp_obj_t value_in) {
+    const char *key = mp_obj_str_get_str(key_in);
+    const char *value = mp_obj_str_get_str(value_in);
+    int ret;
+
+    #if _WIN32
+    ret = _putenv_s(key, value);
+    #else
+    ret = setenv(key, value, 1);
+    #endif
+
+    if (ret == -1) {
+        mp_raise_OSError(errno);
+    }
+    return mp_const_none;
+}
+MP_DEFINE_CONST_FUN_OBJ_2(mod_os_putenv_obj, mod_os_putenv);
+
+STATIC mp_obj_t mod_os_unsetenv(mp_obj_t key_in) {
+    const char *key = mp_obj_str_get_str(key_in);
+    int ret;
+
+    #if _WIN32
+    ret = _putenv_s(key, "");
+    #else
+    ret = unsetenv(key);
+    #endif
+
+    if (ret == -1) {
+        mp_raise_OSError(errno);
+    }
+    return mp_const_none;
+}
+MP_DEFINE_CONST_FUN_OBJ_1(mod_os_unsetenv_obj, mod_os_unsetenv);
+
 STATIC mp_obj_t mod_os_mkdir(mp_obj_t path_in) {
     // TODO: Accept mode param
     const char *path = mp_obj_str_get_str(path_in);
@@ -283,6 +318,8 @@ STATIC const mp_rom_map_elem_t mp_module_os_globals_table[] = {
     { MP_ROM_QSTR(MP_QSTR_rename), MP_ROM_PTR(&mod_os_rename_obj) },
     { MP_ROM_QSTR(MP_QSTR_rmdir), MP_ROM_PTR(&mod_os_rmdir_obj) },
     { MP_ROM_QSTR(MP_QSTR_getenv), MP_ROM_PTR(&mod_os_getenv_obj) },
+    { MP_ROM_QSTR(MP_QSTR_putenv), MP_ROM_PTR(&mod_os_putenv_obj) },
+    { MP_ROM_QSTR(MP_QSTR_unsetenv), MP_ROM_PTR(&mod_os_unsetenv_obj) },
     { MP_ROM_QSTR(MP_QSTR_mkdir), MP_ROM_PTR(&mod_os_mkdir_obj) },
     { MP_ROM_QSTR(MP_QSTR_ilistdir), MP_ROM_PTR(&mod_os_ilistdir_obj) },
     #if MICROPY_PY_OS_DUPTERM

diff --git a/ports/unix/moduos_vfs.c b/ports/unix/moduos_vfs.c
@@ -37,6 +37,8 @@
 // These are defined in modos.c
 MP_DECLARE_CONST_FUN_OBJ_VAR_BETWEEN(mod_os_errno_obj);
 MP_DECLARE_CONST_FUN_OBJ_1(mod_os_getenv_obj);
+MP_DECLARE_CONST_FUN_OBJ_1(mod_os_putenv_obj);
+MP_DECLARE_CONST_FUN_OBJ_1(mod_os_unsetenv_obj);
 MP_DECLARE_CONST_FUN_OBJ_1(mod_os_system_obj);
 
 STATIC const mp_rom_map_elem_t uos_vfs_module_globals_table[] = {
@@ -45,6 +47,8 @@ STATIC const mp_rom_map_elem_t uos_vfs_module_globals_table[] = {
 
     { MP_ROM_QSTR(MP_QSTR_errno), MP_ROM_PTR(&mod_os_errno_obj) },
     { MP_ROM_QSTR(MP_QSTR_getenv), MP_ROM_PTR(&mod_os_getenv_obj) },
+    { MP_ROM_QSTR(MP_QSTR_putenv), MP_ROM_PTR(&mod_os_putenv_obj) },
+    { MP_ROM_QSTR(MP_QSTR_unsetenv), MP_ROM_PTR(&mod_os_unsetenv_obj) },
     { MP_ROM_QSTR(MP_QSTR_system), MP_ROM_PTR(&mod_os_system_obj) },
 
     { MP_ROM_QSTR(MP_QSTR_mount), MP_ROM_PTR(&mp_vfs_mount_obj) },