bpo-39413: Implement os.unsetenv() on Windows (GH-18163)

vstinner · web-flow · commit 161e7b36b1ea · 2020-01-24T11:53:44.000+01:00
The os.unsetenv() function is now also available on Windows.
diff --git a/Doc/library/os.rst b/Doc/library/os.rst
@@ -645,6 +645,9 @@ process and user.
 
    .. availability:: most flavors of Unix, Windows.
 
+   .. versionchanged:: 3.9
+      The function is now also available on Windows.
+
 
 .. _os-newstreams:
 
diff --git a/Doc/whatsnew/3.9.rst b/Doc/whatsnew/3.9.rst
@@ -224,6 +224,9 @@ Exposed the Linux-specific :func:`os.pidfd_open` (:issue:`38692`) and
 :data:`os.P_PIDFD` (:issue:`38713`) for process management with file
 descriptors.
 
+The :func:`os.unsetenv` function is now also available on Windows.
+(Contributed by Victor Stinner in :issue:`39413`.)
+
 poplib
 ------
 
diff --git a/Lib/test/test_os.py b/Lib/test/test_os.py
@@ -953,17 +953,44 @@ def test_environb(self):
         value_str = value.decode(sys.getfilesystemencoding(), 'surrogateescape')
         self.assertEqual(os.environ['bytes'], value_str)
 
+    @unittest.skipUnless(hasattr(os, 'putenv'), "Test needs os.putenv()")
+    @unittest.skipUnless(hasattr(os, 'unsetenv'), "Test needs os.unsetenv()")
+    def test_putenv_unsetenv(self):
+        name = "PYTHONTESTVAR"
+        value = "testvalue"
+        code = f'import os; print(repr(os.environ.get({name!r})))'
+
+        with support.EnvironmentVarGuard() as env:
+            env.pop(name, None)
+
+            os.putenv(name, value)
+            proc = subprocess.run([sys.executable, '-c', code], check=True,
+                                  stdout=subprocess.PIPE, text=True)
+            self.assertEqual(proc.stdout.rstrip(), repr(value))
+
+            os.unsetenv(name)
+            proc = subprocess.run([sys.executable, '-c', code], check=True,
+                                  stdout=subprocess.PIPE, text=True)
+            self.assertEqual(proc.stdout.rstrip(), repr(None))
+
     # On OS X < 10.6, unsetenv() doesn't return a value (bpo-13415).
     @support.requires_mac_ver(10, 6)
-    def test_unset_error(self):
+    @unittest.skipUnless(hasattr(os, 'putenv'), "Test needs os.putenv()")
+    @unittest.skipUnless(hasattr(os, 'unsetenv'), "Test needs os.unsetenv()")
+    def test_putenv_unsetenv_error(self):
+        # Empty variable name is invalid.
+        # "=" and null character are not allowed in a variable name.
+        for name in ('', '=name', 'na=me', 'name=', 'name\0', 'na\0me'):
+            self.assertRaises((OSError, ValueError), os.putenv, name, "value")
+            self.assertRaises((OSError, ValueError), os.unsetenv, name)
+
         if sys.platform == "win32":
-            # an environment variable is limited to 32,767 characters
-            key = 'x' * 50000
-            self.assertRaises(ValueError, os.environ.__delitem__, key)
-        else:
-            # "=" is not allowed in a variable name
-            key = 'key='
-            self.assertRaises(OSError, os.environ.__delitem__, key)
+            # On Windows, an environment variable string ("name=value" string)
+            # is limited to 32,767 characters
+            longstr = 'x' * 32_768
+            self.assertRaises(ValueError, os.putenv, longstr, "1")
+            self.assertRaises(ValueError, os.putenv, "X", longstr)
+            self.assertRaises(ValueError, os.unsetenv, longstr)
 
     def test_key_type(self):
         missing = 'missingkey'
diff --git a/Misc/NEWS.d/next/Library/2020-01-24-10-10-25.bpo-39413.7XYDM8.rst b/Misc/NEWS.d/next/Library/2020-01-24-10-10-25.bpo-39413.7XYDM8.rst
@@ -0,0 +1 @@
+The :func:`os.unsetenv` function is now also available on Windows.
diff --git a/Modules/clinic/posixmodule.c.h b/Modules/clinic/posixmodule.c.h
diff --git a/Modules/posixmodule.c b/Modules/posixmodule.c
@@ -10083,23 +10083,9 @@ posix_putenv_dict_setitem(PyObject *name, PyObject *value)
 
 
 #ifdef MS_WINDOWS
-/*[clinic input]
-os.putenv
-
-    name: unicode
-    value: unicode
-    /
-
-Change or add an environment variable.
-[clinic start generated code]*/
-
-static PyObject *
-os_putenv_impl(PyObject *module, PyObject *name, PyObject *value)
-/*[clinic end generated code: output=d29a567d6b2327d2 input=ba586581c2e6105f]*/
+static PyObject*
+win32_putenv(PyObject *name, PyObject *value)
 {
-    const wchar_t *env;
-    Py_ssize_t size;
-
     /* Search from index 1 because on Windows starting '=' is allowed for
        defining hidden environment variables. */
     if (PyUnicode_GET_LENGTH(name) == 0 ||
@@ -10108,36 +10094,68 @@ os_putenv_impl(PyObject *module, PyObject *name, PyObject *value)
         PyErr_SetString(PyExc_ValueError, "illegal environment variable name");
         return NULL;
     }
-    PyObject *unicode = PyUnicode_FromFormat("%U=%U", name, value);
+    PyObject *unicode;
+    if (value != NULL) {
+        unicode = PyUnicode_FromFormat("%U=%U", name, value);
+    }
+    else {
+        unicode = PyUnicode_FromFormat("%U=", name);
+    }
     if (unicode == NULL) {
         return NULL;
     }
 
-    env = PyUnicode_AsUnicodeAndSize(unicode, &size);
-    if (env == NULL)
-        goto error;
+    Py_ssize_t size;
+    /* PyUnicode_AsWideCharString() rejects embedded null characters */
+    wchar_t *env = PyUnicode_AsWideCharString(unicode, &size);
+    Py_DECREF(unicode);
+
+    if (env == NULL) {
+        return NULL;
+    }
     if (size > _MAX_ENV) {
         PyErr_Format(PyExc_ValueError,
                      "the environment variable is longer than %u characters",
                      _MAX_ENV);
-        goto error;
-    }
-    if (wcslen(env) != (size_t)size) {
-        PyErr_SetString(PyExc_ValueError, "embedded null character");
-        goto error;
+        PyMem_Free(env);
+        return NULL;
     }
 
-    if (_wputenv(env)) {
+    /* _wputenv() and SetEnvironmentVariableW() update the environment in the
+       Process Environment Block (PEB). _wputenv() also updates CRT 'environ'
+       and '_wenviron' variables, whereas SetEnvironmentVariableW() does not.
+
+       Prefer _wputenv() to be compatible with C libraries using CRT
+       variables and CRT functions using these variables (ex: getenv()). */
+    int err = _wputenv(env);
+    PyMem_Free(env);
+
+    if (err) {
         posix_error();
-        goto error;
+        return NULL;
     }
-    Py_DECREF(unicode);
 
     Py_RETURN_NONE;
+}
+#endif
 
-error:
-    Py_DECREF(unicode);
-    return NULL;
+
+#ifdef MS_WINDOWS
+/*[clinic input]
+os.putenv
+
+    name: unicode
+    value: unicode
+    /
+
+Change or add an environment variable.
+[clinic start generated code]*/
+
+static PyObject *
+os_putenv_impl(PyObject *module, PyObject *name, PyObject *value)
+/*[clinic end generated code: output=d29a567d6b2327d2 input=ba586581c2e6105f]*/
+{
+    return win32_putenv(name, value);
 }
 /* repeat !defined(MS_WINDOWS) to workaround an Argument Clinic issue */
 #elif (defined(HAVE_SETENV) || defined(HAVE_PUTENV)) && !defined(MS_WINDOWS)
@@ -10186,7 +10204,23 @@ os_putenv_impl(PyObject *module, PyObject *name, PyObject *value)
 #endif  /* defined(HAVE_SETENV) || defined(HAVE_PUTENV) */
 
 
-#ifdef HAVE_UNSETENV
+#ifdef MS_WINDOWS
+/*[clinic input]
+os.unsetenv
+    name: unicode
+    /
+
+Delete an environment variable.
+[clinic start generated code]*/
+
+static PyObject *
+os_unsetenv_impl(PyObject *module, PyObject *name)
+/*[clinic end generated code: output=54c4137ab1834f02 input=4d6a1747cc526d2f]*/
+{
+    return win32_putenv(name, NULL);
+}
+/* repeat !defined(MS_WINDOWS) to workaround an Argument Clinic issue */
+#elif defined(HAVE_UNSETENV) && !defined(MS_WINDOWS)
 /*[clinic input]
 os.unsetenv
     name: FSConverter
@@ -10199,16 +10233,13 @@ static PyObject *
 os_unsetenv_impl(PyObject *module, PyObject *name)
 /*[clinic end generated code: output=54c4137ab1834f02 input=2bb5288a599c7107]*/
 {
-#ifndef HAVE_BROKEN_UNSETENV
-    int err;
-#endif
-
 #ifdef HAVE_BROKEN_UNSETENV
     unsetenv(PyBytes_AS_STRING(name));
 #else
-    err = unsetenv(PyBytes_AS_STRING(name));
-    if (err)
+    int err = unsetenv(PyBytes_AS_STRING(name));
+    if (err) {
         return posix_error();
+    }
 #endif
 
 #ifdef PY_PUTENV_DICT

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	+The :func:`os.unsetenv` function is now also available on Windows.