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

Author: ambv

Doc/c-api/init_config.rst

@@ -716,7 +716,7 @@ PyConfig
 
       Set to ``1`` by the :envvar:`PYTHONDUMPREFS` environment variable.
 
-      Need a special build of Python with the ``Py_TRACE_REFS`` macro defined:
+      Needs a special build of Python with the ``Py_TRACE_REFS`` macro defined:
       see the :option:`configure --with-trace-refs option <--with-trace-refs>`.
 
       Default: ``0``.
@@ -1048,7 +1048,7 @@ PyConfig
       Incremented by the :option:`-d` command line option. Set to the
       :envvar:`PYTHONDEBUG` environment variable value.
 
-      Need a :ref:`debug build of Python <debug-build>` (the ``Py_DEBUG`` macro
+      Needs a :ref:`debug build of Python <debug-build>` (the ``Py_DEBUG`` macro
       must be defined).
 
       Default: ``0``.
@@ -1100,6 +1100,7 @@ PyConfig
 
       Set by the :option:`-X pycache_prefix=PATH <-X>` command line option and
       the :envvar:`PYTHONPYCACHEPREFIX` environment variable.
+      The command-line option takes precedence.
 
       If ``NULL``, :data:`sys.pycache_prefix` is set to ``None``.
 
@@ -1143,13 +1144,27 @@ PyConfig
 
       Default: ``NULL``.
 
+   .. c:member:: wchar_t* run_presite
+
+      ``package.module`` path to module that should be imported before
+      ``site.py`` is run.
+
+      Set by the :option:`-X presite=package.module <-X>` command-line
+      option and the :envvar:`PYTHON_PRESITE` environment variable.
+      The command-line option takes precedence.
+
+      Needs a :ref:`debug build of Python <debug-build>` (the ``Py_DEBUG`` macro
+      must be defined).
+
+      Default: ``NULL``.
+
    .. c:member:: int show_ref_count
 
       Show total reference count at exit (excluding immortal objects)?
 
       Set to ``1`` by :option:`-X showrefcount <-X>` command line option.
 
-      Need a :ref:`debug build of Python <debug-build>` (the ``Py_REF_DEBUG``
+      Needs a :ref:`debug build of Python <debug-build>` (the ``Py_REF_DEBUG``
       macro must be defined).
 
       Default: ``0``.

Doc/using/cmdline.rst

@@ -552,6 +552,12 @@ Miscellaneous options
      This option may be useful for users who need to limit CPU resources of a
      container system. See also :envvar:`PYTHON_CPU_COUNT`.
      If *n* is ``default``, nothing is overridden.
+   * :samp:`-X presite={package.module}` specifies a module that should be
+     imported before the :mod:`site` module is executed and before the
+     :mod:`__main__` module exists.  Therefore, the imported module isn't
+     :mod:`__main__`. This can be used to execute code early during Python
+     initialization. Python needs to be :ref:`built in debug mode <debug-build>`
+     for this option to exist.  See also :envvar:`PYTHON_PRESITE`.
 
    It also allows passing arbitrary values and retrieving them through the
    :data:`sys._xoptions` dictionary.
@@ -602,6 +608,9 @@ Miscellaneous options
    .. versionadded:: 3.13
       The ``-X cpu_count`` option.
 
+   .. versionadded:: 3.13
+      The ``-X presite`` option.
+
 
 Options you shouldn't use
 ~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -1091,13 +1100,33 @@ Debug-mode variables
    If set, Python will dump objects and reference counts still alive after
    shutting down the interpreter.
 
-   Need Python configured with the :option:`--with-trace-refs` build option.
+   Needs Python configured with the :option:`--with-trace-refs` build option.
 
-.. envvar:: PYTHONDUMPREFSFILE=FILENAME
+.. envvar:: PYTHONDUMPREFSFILE
 
    If set, Python will dump objects and reference counts still alive
-   after shutting down the interpreter into a file called *FILENAME*.
+   after shutting down the interpreter into a file under the path given
+   as the value to this environment variable.
 
-   Need Python configured with the :option:`--with-trace-refs` build option.
+   Needs Python configured with the :option:`--with-trace-refs` build option.
 
    .. versionadded:: 3.11
+
+.. envvar:: PYTHON_PRESITE
+
+   If this variable is set to a module, that module will be imported
+   early in the interpreter lifecycle, before the :mod:`site` module is
+   executed, and before the :mod:`__main__` module is created.
+   Therefore, the imported module is not treated as :mod:`__main__`.
+
+   This can be used to execute code early during Python initialization.
+
+   To import a submodule, use ``package.module`` as the value, like in
+   an import statement.
+
+   See also the :option:`-X presite <-X>` command-line option,
+   which takes precedence over this variable.
+
+   Needs Python configured with the :option:`--with-pydebug` build option.
+
+   .. versionadded:: 3.13

Doc/whatsnew/3.13.rst

@@ -1356,3 +1356,13 @@ removed, although there is currently no date scheduled for their removal.
 
 * Remove undocumented ``PY_TIMEOUT_MAX`` constant from the limited C API.
   (Contributed by Victor Stinner in :gh:`110014`.)
+
+
+Regression Test Changes
+=======================
+
+* Python built with :file:`configure` :option:`--with-pydebug` now
+  supports a :option:`-X presite=package.module <-X>` command-line
+  option. If used, it specifies a module that should be imported early
+  in the lifecycle of the interpreter, before ``site.py`` is executed.
+  (Contributed by Łukasz Langa in :gh:`110769`.)

Include/cpython/initconfig.h

@@ -225,6 +225,12 @@ typedef struct PyConfig {
     // If non-zero, turns on statistics gathering.
     int _pystats;
 #endif
+
+#ifdef Py_DEBUG
+    // If not empty, import a non-__main__ module before site.py is executed.
+    // PYTHON_PRESITE=package.module or -X presite=package.module
+    wchar_t *run_presite;
+#endif
 } PyConfig;
 
 PyAPI_FUNC(void) PyConfig_InitPythonConfig(PyConfig *config);

Lib/test/test_embed.py

@@ -515,6 +515,8 @@ class InitConfigTests(EmbeddingTestsMixin, unittest.TestCase):
     }
     if Py_STATS:
         CONFIG_COMPAT['_pystats'] = 0
+    if support.Py_DEBUG:
+        CONFIG_COMPAT['run_presite'] = None
     if MS_WINDOWS:
         CONFIG_COMPAT.update({
             'legacy_windows_stdio': 0,
@@ -1818,6 +1820,22 @@ def test_no_memleak(self):
                 self.assertEqual(refs, 0, out)
                 self.assertEqual(blocks, 0, out)
 
+    @unittest.skipUnless(support.Py_DEBUG,
+                         '-X presite requires a Python debug build')
+    def test_presite(self):
+        cmd = [sys.executable, "-I", "-X", "presite=test.reperf", "-c", "print('cmd')"]
+        proc = subprocess.run(
+            cmd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            text=True,
+        )
+        self.assertEqual(proc.returncode, 0)
+        out = proc.stdout.strip()
+        self.assertIn("10 times sub", out)
+        self.assertIn("CPU seconds", out)
+        self.assertIn("cmd", out)
+
 
 class StdPrinterTests(EmbeddingTestsMixin, unittest.TestCase):
     # Test PyStdPrinter_Type which is used by _PySys_SetPreliminaryStderr():

Misc/NEWS.d/next/Core and Builtins/2023-10-12-17-15-23.gh-issue-110722.sjMwQe.rst

@@ -0,0 +1,3 @@
+Add :envvar:`PYTHON_PRESITE=package.module` to import a module early in the
+interpreter lifecycle before ``site.py`` is executed.  Python needs to be
+:ref:`built in debug mode <debug-build>` for this option to exist.

Python/initconfig.c

@@ -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();
 }
 

Python/pylifecycle.c

@@ -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;