bpo-41094: Additional fix for PYTHONSTARTUP. (GH-21119)

python · Jun 24, 2020 · a7dc714 · a7dc714
1 parent 33b79b1
commit a7dc714
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 11 deletions.
diff --git a/Lib/test/test_embed.py b/Lib/test/test_embed.py
@@ -1349,7 +1349,7 @@ def test_audit_run_file(self):
                                       returncode=1)
 
     def test_audit_run_interactivehook(self):
-        startup = os.path.join(self.oldcwd, support.TESTFN) + (support.TESTFN or '') + ".py"
+        startup = os.path.join(self.oldcwd, support.TESTFN) + (support.FS_NONASCII or '') + ".py"
         with open(startup, "w", encoding="utf-8") as f:
             print("import sys", file=f)
             print("sys.__interactivehook__ = lambda: None", file=f)
@@ -1362,7 +1362,7 @@ def test_audit_run_interactivehook(self):
             os.unlink(startup)
 
     def test_audit_run_startup(self):
-        startup = os.path.join(self.oldcwd, support.TESTFN) + (support.TESTFN or '') + ".py"
+        startup = os.path.join(self.oldcwd, support.TESTFN) + (support.FS_NONASCII or '') + ".py"
         with open(startup, "w", encoding="utf-8") as f:
             print("pass", file=f)
         try:

diff --git a/Modules/main.c b/Modules/main.c
@@ -375,36 +375,70 @@ pymain_run_file(const PyConfig *config, PyCompilerFlags *cf)
 static int
 pymain_run_startup(PyConfig *config, PyCompilerFlags *cf, int *exitcode)
 {
+    int ret;
+    PyObject *startup_obj = NULL;
+    if (!config->use_environment) {
+        return 0;
+    }
+#ifdef MS_WINDOWS
+    const wchar_t *wstartup = _wgetenv(L"PYTHONSTARTUP");
+    if (wstartup == NULL || wstartup[0] == L'\0') {
+        return 0;
+    }
+    PyObject *startup_bytes = NULL;
+    startup_obj = PyUnicode_FromWideChar(wstartup, wcslen(wstartup));
+    if (startup_obj == NULL) {
+        goto error;
+    }
+    startup_bytes = PyUnicode_EncodeFSDefault(startup_obj);
+    if (startup_bytes == NULL) {
+        goto error;
+    }
+    const char *startup = PyBytes_AS_STRING(startup_bytes);
+#else
     const char *startup = _Py_GetEnv(config->use_environment, "PYTHONSTARTUP");
     if (startup == NULL) {
         return 0;
     }
-    PyObject *startup_obj = PyUnicode_DecodeFSDefault(startup);
+    startup_obj = PyUnicode_DecodeFSDefault(startup);
     if (startup_obj == NULL) {
-        return pymain_err_print(exitcode);
+        goto error;
     }
+#endif
     if (PySys_Audit("cpython.run_startup", "O", startup_obj) < 0) {
-        Py_DECREF(startup_obj);
-        return pymain_err_print(exitcode);
+        goto error;
     }
-    Py_DECREF(startup_obj);
 
+#ifdef MS_WINDOWS
+    FILE *fp = _Py_wfopen(wstartup, L"r");
+#else
     FILE *fp = _Py_fopen(startup, "r");
+#endif
     if (fp == NULL) {
         int save_errno = errno;
         PyErr_Clear();
         PySys_WriteStderr("Could not open PYTHONSTARTUP\n");
 
         errno = save_errno;
-        PyErr_SetFromErrnoWithFilename(PyExc_OSError, startup);
-
-        return pymain_err_print(exitcode);
+        PyErr_SetFromErrnoWithFilenameObjects(PyExc_OSError, startup_obj, NULL);
+        goto error;
     }
 
     (void) PyRun_SimpleFileExFlags(fp, startup, 0, cf);
     PyErr_Clear();
     fclose(fp);
-    return 0;
+    ret = 0;
+
+done:
+#ifdef MS_WINDOWS
+    Py_XDECREF(startup_bytes);
+#endif
+    Py_XDECREF(startup_obj);
+    return ret;
+
+error:
+    ret = pymain_err_print(exitcode);
+    goto done;
 }