bpo-36142: Add _PyPreConfig_SetAllocator() (GH-12187) · python/cpython@c656e25
@@ -125,15 +125,8 @@ precmdline_clear(_PyPreCmdline *cmdline)
125125void
126126_PyPreConfig_Clear(_PyPreConfig *config)
127127{
128-#define CLEAR(ATTR) \
129- do { \
130- PyMem_RawFree(ATTR); \
131- ATTR = NULL; \
132- } while (0)
133-134-CLEAR(config->allocator);
135-136-#undef CLEAR
128+PyMem_RawFree(config->allocator);
129+config->allocator = NULL;
137130}
138131139132@@ -453,8 +446,7 @@ preconfig_read(_PyPreConfig *config, const _PyPreCmdline *cmdline)
453446454447/* allocator */
455448if (config->dev_mode && config->allocator == NULL) {
456-const char *allocator = _PyMem_GetDebugAllocatorsName();
457-config->allocator = _PyMem_RawStrdup(allocator);
449+config->allocator = _PyMem_RawStrdup("debug");
458450if (config->allocator == NULL) {
459451return _Py_INIT_NO_MEMORY();
460452 }
@@ -742,31 +734,56 @@ _PyPreConfig_ReadFromArgv(_PyPreConfig *config, const _PyArgv *args)
742734743735744736static _PyInitError
745-_PyPreConfig_Reconfigure(const _PyPreConfig *config)
737+_PyPreConfig_SetAllocator(_PyPreConfig *config)
746738{
747-if (config->allocator != NULL) {
748-const char *allocator = _PyMem_GetAllocatorsName();
749-if (allocator == NULL || strcmp(config->allocator, allocator) != 0) {
750-return _Py_INIT_USER_ERR("cannot modify memory allocator "
751-"after first Py_Initialize()");
752- }
739+assert(!_PyRuntime.core_initialized);
740+741+PyMemAllocatorEx old_alloc;
742+PyMem_GetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
743+744+if (_PyMem_SetupAllocators(config->allocator) < 0) {
745+return _Py_INIT_USER_ERR("Unknown PYTHONMALLOC allocator");
746+ }
747+748+/* Copy the pre-configuration with the new allocator */
749+_PyPreConfig config2 = _PyPreConfig_INIT;
750+if (_PyPreConfig_Copy(&config2, config) < 0) {
751+_PyPreConfig_Clear(&config2);
752+PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
753+return _Py_INIT_NO_MEMORY();
753754 }
755+756+/* Free the old config and replace config with config2. Since config now
757+ owns the data, don't free config2. */
758+PyMemAllocatorEx new_alloc;
759+PyMem_GetAllocator(PYMEM_DOMAIN_RAW, &new_alloc);
760+PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
761+_PyPreConfig_Clear(config);
762+PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &new_alloc);
763+764+*config = config2;
765+754766return _Py_INIT_OK();
755767}
756768757769770+/* Write the pre-configuration.
771+772+ If the memory allocator is changed, config is re-allocated with new
773+ allocator. So calling _PyPreConfig_Clear(config) is safe after this call. */
758774_PyInitError
759-_PyPreConfig_Write(const _PyPreConfig *config)
775+_PyPreConfig_Write(_PyPreConfig *config)
760776{
761777if (_PyRuntime.core_initialized) {
762778/* bpo-34008: Calling Py_Main() after Py_Initialize() ignores
763779 the new configuration. */
764-return _PyPreConfig_Reconfigure(config);
780+return _Py_INIT_OK();
765781 }
766782767783if (config->allocator != NULL) {
768-if (_PyMem_SetupAllocators(config->allocator) < 0) {
769-return _Py_INIT_USER_ERR("Unknown PYTHONMALLOC allocator");
784+_PyInitError err = _PyPreConfig_SetAllocator(config);
785+if (_Py_INIT_FAILED(err)) {
786+return err;
770787 }
771788 }
772789