bpo-36142: _PyPreConfig_Write() sets the allocator (GH-12186)

* _PyPreConfig_Write() now sets the memory allocator.
* _PyPreConfig_Write() gets a return type: _PyInitError.
* _Py_InitializeCore() now reads and writes the pre-configuration
  (set the memory allocator, configure the locale) before reading and
  writing the core configuration.

1 files changed, 29 insertions, 1 deletions

diff --git a/Python/preconfig.c b/Python/preconfig.c
index 46e1809..6924203 100644
--- a/Python/preconfig.c
+++ b/Python/preconfig.c
@@ -741,9 +741,35 @@ done:
 }
 
 
-void
+static _PyInitError
+_PyPreConfig_Reconfigure(const _PyPreConfig *config)
+{
+    if (config->allocator != NULL) {
+        const char *allocator = _PyMem_GetAllocatorsName();
+        if (allocator == NULL || strcmp(config->allocator, allocator) != 0) {
+            return _Py_INIT_USER_ERR("cannot modify memory allocator "
+                                     "after first Py_Initialize()");
+        }
+    }
+    return _Py_INIT_OK();
+}
+
+
+_PyInitError
 _PyPreConfig_Write(const _PyPreConfig *config)
 {
+    if (_PyRuntime.core_initialized) {
+        /* bpo-34008: Calling Py_Main() after Py_Initialize() ignores
+           the new configuration. */
+        return _PyPreConfig_Reconfigure(config);
+    }
+
+    if (config->allocator != NULL) {
+        if (_PyMem_SetupAllocators(config->allocator) < 0) {
+            return _Py_INIT_USER_ERR("Unknown PYTHONMALLOC allocator");
+        }
+    }
+
     _PyPreConfig_SetGlobalConfig(config);
 
     if (config->coerce_c_locale) {
@@ -752,4 +778,6 @@ _PyPreConfig_Write(const _PyPreConfig *config)
 
     /* Set LC_CTYPE to the user preferred locale */
     _Py_SetLocaleFromEnv(LC_CTYPE);
+
+    return _Py_INIT_OK();
 }