gh-110722: Add PYTHON_PRESITE to import a module before site.py is run (#110769)

This commit is contained in:
Łukasz Langa 2023-10-14 23:32:57 +02:00 committed by GitHub
parent ab08ff7882
commit 84b7e9e3fa
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 178 additions and 7 deletions

View file

@ -117,6 +117,9 @@ static const PyConfigSpec PYCONFIG_SPEC[] = {
SPEC(_is_python_build, UINT),
#ifdef Py_STATS
SPEC(_pystats, UINT),
#endif
#ifdef Py_DEBUG
SPEC(run_presite, WSTR_OPT),
#endif
{NULL, 0, 0},
};
@ -241,6 +244,11 @@ The following implementation-specific options are available:\n\
\n\
-X pystats: Enable pystats collection at startup."
#endif
#ifdef Py_DEBUG
"\n\
\n\
-X presite=package.module: import this module before site.py is run."
#endif
;
/* Envvars that don't have equivalent command-line options are listed first */
@ -297,6 +305,9 @@ static const char usage_envvars[] =
#ifdef Py_STATS
"PYTHONSTATS : turns on statistics gathering\n"
#endif
#ifdef Py_DEBUG
"PYTHON_PRESITE=pkg.mod : import this module before site.py is run\n"
#endif
;
#if defined(MS_WINDOWS)
@ -790,6 +801,9 @@ PyConfig_Clear(PyConfig *config)
CLEAR(config->run_module);
CLEAR(config->run_filename);
CLEAR(config->check_hash_pycs_mode);
#ifdef Py_DEBUG
CLEAR(config->run_presite);
#endif
_PyWideStringList_Clear(&config->orig_argv);
#undef CLEAR
@ -1806,6 +1820,36 @@ config_init_pycache_prefix(PyConfig *config)
}
#ifdef Py_DEBUG
static PyStatus
config_init_run_presite(PyConfig *config)
{
assert(config->run_presite == NULL);
const wchar_t *xoption = config_get_xoption(config, L"presite");
if (xoption) {
const wchar_t *sep = wcschr(xoption, L'=');
if (sep && wcslen(sep) > 1) {
config->run_presite = _PyMem_RawWcsdup(sep + 1);
if (config->run_presite == NULL) {
return _PyStatus_NO_MEMORY();
}
}
else {
// PYTHON_PRESITE env var ignored
// if "-X presite=" option is used
config->run_presite = NULL;
}
return _PyStatus_OK();
}
return CONFIG_GET_ENV_DUP(config, &config->run_presite,
L"PYTHON_PRESITE",
"PYTHON_PRESITE");
}
#endif
static PyStatus
config_read_complex_options(PyConfig *config)
{
@ -1861,6 +1905,16 @@ config_read_complex_options(PyConfig *config)
return status;
}
}
#ifdef Py_DEBUG
if (config->run_presite == NULL) {
status = config_init_run_presite(config);
if (_PyStatus_EXCEPTION(status)) {
return status;
}
}
#endif
return _PyStatus_OK();
}

View file

@ -1076,6 +1076,38 @@ pyinit_main_reconfigure(PyThreadState *tstate)
}
#ifdef Py_DEBUG
static void
run_presite(PyThreadState *tstate)
{
PyInterpreterState *interp = tstate->interp;
const PyConfig *config = _PyInterpreterState_GetConfig(interp);
if (!config->run_presite) {
return;
}
PyObject *presite_modname = PyUnicode_FromWideChar(
config->run_presite,
wcslen(config->run_presite)
);
if (presite_modname == NULL) {
fprintf(stderr, "Could not convert pre-site module name to unicode\n");
Py_DECREF(presite_modname);
}
else {
PyObject *presite = PyImport_Import(presite_modname);
if (presite == NULL) {
fprintf(stderr, "pre-site import failed:\n");
_PyErr_Print(tstate);
}
Py_XDECREF(presite);
Py_DECREF(presite_modname);
}
}
#endif
static PyStatus
init_interp_main(PyThreadState *tstate)
{
@ -1157,6 +1189,10 @@ init_interp_main(PyThreadState *tstate)
return status;
}
#ifdef Py_DEBUG
run_presite(tstate);
#endif
status = add_main_module(interp);
if (_PyStatus_EXCEPTION(status)) {
return status;