bpo-45211: Remember the stdlib dir during startup. (gh-28586)

ericsnowcurrently · web-flow · commit 0c50b8c0b827 · 2021-09-28T12:18:28.000-06:00
During runtime startup we figure out the stdlib dir but currently throw that information away. This change preserves it and exposes it via PyConfig.stdlib_dir, _Py_GetStdlibDir(), and sys._stdlib_dir. https://bugs.python.org/issue45211
diff --git a/Include/cpython/initconfig.h b/Include/cpython/initconfig.h
@@ -184,6 +184,7 @@ typedef struct PyConfig {
     /* --- Path configuration outputs ----------- */
     int module_search_paths_set;
     PyWideStringList module_search_paths;
+    wchar_t *stdlib_dir;
     wchar_t *executable;
     wchar_t *base_executable;
     wchar_t *prefix;
diff --git a/Include/internal/pycore_pathconfig.h b/Include/internal/pycore_pathconfig.h
@@ -13,6 +13,7 @@ typedef struct _PyPathConfig {
     wchar_t *program_full_path;
     wchar_t *prefix;
     wchar_t *exec_prefix;
+    wchar_t *stdlib_dir;
     /* Set by Py_SetPath(), or computed by _PyConfig_InitPathConfig() */
     wchar_t *module_search_path;
     /* Python program name */
diff --git a/Include/internal/pycore_pylifecycle.h b/Include/internal/pycore_pylifecycle.h
@@ -122,6 +122,7 @@ PyAPI_FUNC(PyStatus) _Py_PreInitializeFromConfig(
     const PyConfig *config,
     const struct _PyArgv *args);
 
+PyAPI_FUNC(wchar_t *) _Py_GetStdlibDir(void);
 
 PyAPI_FUNC(int) _Py_HandleSystemExit(int *exitcode_p);
 
diff --git a/Lib/test/test_embed.py b/Lib/test/test_embed.py
@@ -406,6 +406,7 @@ class InitConfigTests(EmbeddingTestsMixin, unittest.TestCase):
         'module_search_paths': GET_DEFAULT_CONFIG,
         'module_search_paths_set': 1,
         'platlibdir': sys.platlibdir,
+        'stdlib_dir': GET_DEFAULT_CONFIG,
 
         'site_import': 1,
         'bytes_warning': 0,
@@ -515,6 +516,7 @@ class InitConfigTests(EmbeddingTestsMixin, unittest.TestCase):
         'exec_prefix',
         'program_name',
         'home',
+        'stdlib_dir',
         # program_full_path and module_search_path are copied indirectly from
         # the core configuration in check_path_config().
     ]
@@ -1142,6 +1144,9 @@ def test_init_setpath(self):
             'base_prefix': '',
             'exec_prefix': '',
             'base_exec_prefix': '',
+             # The current getpath.c doesn't determine the stdlib dir
+             # in this case.
+            'stdlib_dir': '',
         }
         self.default_program_name(config)
         env = {'TESTPATH': os.path.pathsep.join(paths)}
@@ -1162,6 +1167,9 @@ def test_init_setpath_config(self):
             'base_prefix': '',
             'exec_prefix': '',
             'base_exec_prefix': '',
+             # The current getpath.c doesn't determine the stdlib dir
+             # in this case.
+            'stdlib_dir': '',
             # overriden by PyConfig
             'program_name': 'conf_program_name',
             'base_executable': 'conf_executable',
@@ -1251,6 +1259,7 @@ def test_init_setpythonhome(self):
             'exec_prefix': exec_prefix,
             'base_exec_prefix': exec_prefix,
             'pythonpath_env': paths_str,
+            'stdlib_dir': home,
         }
         self.default_program_name(config)
         env = {'TESTHOME': home, 'PYTHONPATH': paths_str}
@@ -1288,6 +1297,9 @@ def test_init_pybuilddir(self):
                 'base_executable': executable,
                 'executable': executable,
                 'module_search_paths': module_search_paths,
+                # The current getpath.c doesn't determine the stdlib dir
+                # in this case.
+                'stdlib_dir': None,
             }
             env = self.copy_paths_by_env(config)
             self.check_all_configs("test_init_compat_config", config,
@@ -1345,6 +1357,7 @@ def test_init_pyvenv_cfg(self):
             if MS_WINDOWS:
                 config['base_prefix'] = pyvenv_home
                 config['prefix'] = pyvenv_home
+                config['stdlib_dir'] = os.path.join(pyvenv_home, 'lib')
 
                 ver = sys.version_info
                 dll = f'python{ver.major}'
@@ -1353,6 +1366,10 @@ def test_init_pyvenv_cfg(self):
                 dll += '.DLL'
                 dll = os.path.join(os.path.dirname(executable), dll)
                 path_config['python3_dll'] = dll
+            else:
+                # The current getpath.c doesn't determine the stdlib dir
+                # in this case.
+                config['stdlib_dir'] = None
 
             env = self.copy_paths_by_env(config)
             self.check_all_configs("test_init_compat_config", config,
diff --git a/Lib/test/test_sys.py b/Lib/test/test_sys.py
@@ -13,6 +13,7 @@
 from test.support import os_helper
 from test.support.script_helper import assert_python_ok, assert_python_failure
 from test.support import threading_helper
+from test.support import import_helper
 import textwrap
 import unittest
 import warnings
@@ -994,6 +995,15 @@ def test_module_names(self):
         for name in sys.stdlib_module_names:
             self.assertIsInstance(name, str)
 
+    def test_stdlib_dir(self):
+        os = import_helper.import_fresh_module('os')
+        marker = getattr(os, '__file__', None)
+        if marker and not os.path.exists(marker):
+            marker = None
+        expected = os.path.dirname(marker) if marker else None
+        actual = sys._stdlib_dir
+        self.assertEqual(actual, expected)
+
 
 @test.support.cpython_only
 class UnraisableHookTest(unittest.TestCase):
diff --git a/Modules/getpath.c b/Modules/getpath.c
@@ -1492,6 +1492,16 @@ calculate_path(PyCalculatePath *calculate, _PyPathConfig *pathconfig)
         }
     }
 
+    if (pathconfig->stdlib_dir == NULL) {
+        if (calculate->prefix_found) {
+            /* This must be done *before* calculate_set_prefix() is called. */
+            pathconfig->stdlib_dir = _PyMem_RawWcsdup(calculate->prefix);
+            if (pathconfig->stdlib_dir == NULL) {
+                return _PyStatus_NO_MEMORY();
+            }
+        }
+    }
+
     if (pathconfig->prefix == NULL) {
         status = calculate_set_prefix(calculate, pathconfig);
         if (_PyStatus_EXCEPTION(status)) {
diff --git a/PC/getpathp.c b/PC/getpathp.c
@@ -116,6 +116,8 @@
  * with a semicolon separated path prior to calling Py_Initialize.
  */
 
+#define STDLIB_SUBDIR L"lib"
+
 #define INIT_ERR_BUFFER_OVERFLOW() _PyStatus_ERR("buffer overflow")
 
 
@@ -293,12 +295,12 @@ search_for_prefix(wchar_t *prefix, const wchar_t *argv0_path)
     wcscpy_s(stdlibdir, Py_ARRAY_LENGTH(stdlibdir), prefix);
     /* We initialize with the longest possible path, in case it doesn't fit.
        This also gives us an initial SEP at stdlibdir[wcslen(prefix)]. */
-    join(stdlibdir, L"lib");
+    join(stdlibdir, STDLIB_SUBDIR);
     do {
         assert(stdlibdir[wcslen(prefix)] == SEP);
         /* Due to reduce() and our initial value, this result
            is guaranteed to fit. */
-        wcscpy(&stdlibdir[wcslen(prefix) + 1], L"lib");
+        wcscpy(&stdlibdir[wcslen(prefix) + 1], STDLIB_SUBDIR);
         if (is_stdlibdir(stdlibdir)) {
             return 1;
         }
@@ -1013,6 +1015,12 @@ calculate_path(PyCalculatePath *calculate, _PyPathConfig *pathconfig)
     }
 
 done:
+    if (pathconfig->stdlib_dir == NULL) {
+        pathconfig->stdlib_dir = _Py_join_relfile(prefix, STDLIB_SUBDIR);
+        if (pathconfig->stdlib_dir == NULL) {
+            return _PyStatus_NO_MEMORY();
+        }
+    }
     if (pathconfig->prefix == NULL) {
         pathconfig->prefix = _PyMem_RawWcsdup(prefix);
         if (pathconfig->prefix == NULL) {
diff --git a/Python/initconfig.c b/Python/initconfig.c
@@ -669,6 +669,7 @@ PyConfig_Clear(PyConfig *config)
     _PyWideStringList_Clear(&config->xoptions);
     _PyWideStringList_Clear(&config->module_search_paths);
     config->module_search_paths_set = 0;
+    CLEAR(config->stdlib_dir);
 
     CLEAR(config->executable);
     CLEAR(config->base_executable);
@@ -909,6 +910,7 @@ _PyConfig_Copy(PyConfig *config, const PyConfig *config2)
     COPY_WSTRLIST(xoptions);
     COPY_WSTRLIST(module_search_paths);
     COPY_ATTR(module_search_paths_set);
+    COPY_WSTR_ATTR(stdlib_dir);
 
     COPY_WSTR_ATTR(executable);
     COPY_WSTR_ATTR(base_executable);
@@ -1015,6 +1017,7 @@ _PyConfig_AsDict(const PyConfig *config)
     SET_ITEM_WSTR(home);
     SET_ITEM_INT(module_search_paths_set);
     SET_ITEM_WSTRLIST(module_search_paths);
+    SET_ITEM_WSTR(stdlib_dir);
     SET_ITEM_WSTR(executable);
     SET_ITEM_WSTR(base_executable);
     SET_ITEM_WSTR(prefix);
@@ -1318,6 +1321,7 @@ _PyConfig_FromDict(PyConfig *config, PyObject *dict)
     // Path configuration output
     GET_UINT(module_search_paths_set);
     GET_WSTRLIST(module_search_paths);
+    GET_WSTR_OPT(stdlib_dir);
     GET_WSTR_OPT(executable);
     GET_WSTR_OPT(base_executable);
     GET_WSTR_OPT(prefix);
@@ -3094,6 +3098,7 @@ _Py_DumpPathConfig(PyThreadState *tstate)
     PySys_WriteStderr("  environment = %i\n", config->use_environment);
     PySys_WriteStderr("  user site = %i\n", config->user_site_directory);
     PySys_WriteStderr("  import site = %i\n", config->site_import);
+    DUMP_CONFIG("stdlib dir", stdlib_dir);
 #undef DUMP_CONFIG
 
 #define DUMP_SYS(NAME) \
diff --git a/Python/pathconfig.c b/Python/pathconfig.c
@@ -54,6 +54,7 @@ pathconfig_clear(_PyPathConfig *config)
     CLEAR(config->program_full_path);
     CLEAR(config->prefix);
     CLEAR(config->exec_prefix);
+    CLEAR(config->stdlib_dir);
     CLEAR(config->module_search_path);
     CLEAR(config->program_name);
     CLEAR(config->home);
@@ -83,6 +84,7 @@ pathconfig_copy(_PyPathConfig *config, const _PyPathConfig *config2)
     COPY_ATTR(prefix);
     COPY_ATTR(exec_prefix);
     COPY_ATTR(module_search_path);
+    COPY_ATTR(stdlib_dir);
     COPY_ATTR(program_name);
     COPY_ATTR(home);
 #ifdef MS_WINDOWS
@@ -167,6 +169,7 @@ pathconfig_set_from_config(_PyPathConfig *pathconfig, const PyConfig *config)
     COPY_CONFIG(program_full_path, executable);
     COPY_CONFIG(prefix, prefix);
     COPY_CONFIG(exec_prefix, exec_prefix);
+    COPY_CONFIG(stdlib_dir, stdlib_dir);
     COPY_CONFIG(program_name, program_name);
     COPY_CONFIG(home, home);
 #ifdef MS_WINDOWS
@@ -218,6 +221,7 @@ _PyPathConfig_AsDict(void)
     SET_ITEM_STR(prefix);
     SET_ITEM_STR(exec_prefix);
     SET_ITEM_STR(module_search_path);
+    SET_ITEM_STR(stdlib_dir);
     SET_ITEM_STR(program_name);
     SET_ITEM_STR(home);
 #ifdef MS_WINDOWS
@@ -311,6 +315,7 @@ config_init_module_search_paths(PyConfig *config, _PyPathConfig *pathconfig)
 
    - exec_prefix
    - module_search_path
+   - stdlib_dir
    - prefix
    - program_full_path
 
@@ -401,6 +406,7 @@ config_init_pathconfig(PyConfig *config, int compute_path_config)
     COPY_ATTR(program_full_path, executable);
     COPY_ATTR(prefix, prefix);
     COPY_ATTR(exec_prefix, exec_prefix);
+    COPY_ATTR(stdlib_dir, stdlib_dir);
 
 #undef COPY_ATTR
 
@@ -486,16 +492,25 @@ Py_SetPath(const wchar_t *path)
 
     PyMem_RawFree(_Py_path_config.prefix);
     PyMem_RawFree(_Py_path_config.exec_prefix);
+    PyMem_RawFree(_Py_path_config.stdlib_dir);
     PyMem_RawFree(_Py_path_config.module_search_path);
 
     _Py_path_config.prefix = _PyMem_RawWcsdup(L"");
     _Py_path_config.exec_prefix = _PyMem_RawWcsdup(L"");
+    // XXX Copy this from the new module_search_path?
+    if (_Py_path_config.home != NULL) {
+        _Py_path_config.stdlib_dir = _PyMem_RawWcsdup(_Py_path_config.home);
+    }
+    else {
+        _Py_path_config.stdlib_dir = _PyMem_RawWcsdup(L"");
+    }
     _Py_path_config.module_search_path = _PyMem_RawWcsdup(path);
 
     PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
 
     if (_Py_path_config.prefix == NULL
         || _Py_path_config.exec_prefix == NULL
+        || _Py_path_config.stdlib_dir == NULL
         || _Py_path_config.module_search_path == NULL)
     {
         path_out_of_memory(__func__);
@@ -515,10 +530,13 @@ Py_SetPythonHome(const wchar_t *home)
 
     PyMem_RawFree(_Py_path_config.home);
     _Py_path_config.home = _PyMem_RawWcsdup(home);
+    if (_Py_path_config.home != NULL) {
+        _Py_path_config.stdlib_dir = _PyMem_RawWcsdup(home);
+    }
 
     PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
 
-    if (_Py_path_config.home == NULL) {
+    if (_Py_path_config.home == NULL || _Py_path_config.stdlib_dir == NULL) {
         path_out_of_memory(__func__);
     }
 }
@@ -572,6 +590,17 @@ Py_GetPath(void)
 }
 
 
+wchar_t *
+_Py_GetStdlibDir(void)
+{
+    wchar_t *stdlib_dir = _Py_path_config.stdlib_dir;
+    if (stdlib_dir != NULL && stdlib_dir[0] != L'\0') {
+        return stdlib_dir;
+    }
+    return NULL;
+}
+
+
 wchar_t *
 Py_GetPrefix(void)
 {
diff --git a/Python/sysmodule.c b/Python/sysmodule.c
@@ -2974,6 +2974,14 @@ _PySys_UpdateConfig(PyThreadState *tstate)
 
     SET_SYS("_xoptions", sys_create_xoptions_dict(config));
 
+    const wchar_t *stdlibdir = _Py_GetStdlibDir();
+    if (stdlibdir != NULL) {
+        SET_SYS_FROM_WSTR("_stdlib_dir", stdlibdir);
+    }
+    else {
+        PyDict_SetItemString(sysdict, "_stdlib_dir", Py_None);
+    }
+
 #undef SET_SYS_FROM_WSTR
 #undef COPY_LIST
 #undef COPY_WSTR
