bpo-36763: Add PyMemAllocatorName

Include/cpython/coreconfig.h

@@ -120,7 +120,9 @@ typedef struct {
     int utf8_mode;
 
     int dev_mode;           /* Development mode. PYTHONDEVMODE, -X dev */
-    char *allocator;        /* Memory allocator: PYTHONMALLOC */
+
+    /* Memory allocator: PYTHONMALLOC env var */
+    PyMemAllocatorName allocator;
 } _PyPreConfig;
 
 #ifdef MS_WINDOWS
@@ -137,7 +139,7 @@ typedef struct {
         .isolated = -1, \
         .use_environment = -1, \
         .dev_mode = -1, \
-        .allocator = NULL}
+        .allocator = PYMEM_ALLOCATOR_NOT_SET}
 
 
 /* --- _PyCoreConfig ---------------------------------------------- */

Include/cpython/pymem.h

@@ -11,12 +11,8 @@ PyAPI_FUNC(void *) PyMem_RawCalloc(size_t nelem, size_t elsize);
 PyAPI_FUNC(void *) PyMem_RawRealloc(void *ptr, size_t new_size);
 PyAPI_FUNC(void) PyMem_RawFree(void *ptr);
 
-/* Configure the Python memory allocators. Pass NULL to use default
-   allocators. */
-PyAPI_FUNC(int) _PyMem_SetupAllocators(const char *opt);
-
 /* Try to get the allocators name set by _PyMem_SetupAllocators(). */
-PyAPI_FUNC(const char*) _PyMem_GetAllocatorsName(void);
+PyAPI_FUNC(const char*) _PyMem_GetCurrentAllocatorName(void);
 
 PyAPI_FUNC(void *) PyMem_Calloc(size_t nelem, size_t elsize);
 
@@ -41,6 +37,19 @@ typedef enum {
     PYMEM_DOMAIN_OBJ
 } PyMemAllocatorDomain;
 
+typedef enum {
+    PYMEM_ALLOCATOR_NOT_SET = 0,
+    PYMEM_ALLOCATOR_DEFAULT = 1,
+    PYMEM_ALLOCATOR_DEBUG = 2,
+    PYMEM_ALLOCATOR_MALLOC = 3,
+    PYMEM_ALLOCATOR_MALLOC_DEBUG = 4,
+#ifdef WITH_PYMALLOC
+    PYMEM_ALLOCATOR_PYMALLOC = 5,
+    PYMEM_ALLOCATOR_PYMALLOC_DEBUG = 6,
+#endif
+} PyMemAllocatorName;
+
+
 typedef struct {
     /* user context passed as the first argument to the 4 functions */
     void *ctx;

Include/internal/pycore_coreconfig.h

@@ -88,15 +88,14 @@ PyAPI_FUNC(_PyInitError) _PyPreCmdline_Read(_PyPreCmdline *cmdline,
 
 /* --- _PyPreConfig ----------------------------------------------- */
 
-PyAPI_FUNC(void) _PyPreConfig_Clear(_PyPreConfig *config);
 PyAPI_FUNC(int) _PyPreConfig_Copy(_PyPreConfig *config,
     const _PyPreConfig *config2);
 PyAPI_FUNC(PyObject*) _PyPreConfig_AsDict(const _PyPreConfig *config);
 PyAPI_FUNC(void) _PyCoreConfig_GetCoreConfig(_PyPreConfig *config,
     const _PyCoreConfig *core_config);
 PyAPI_FUNC(_PyInitError) _PyPreConfig_Read(_PyPreConfig *config,
     const _PyArgv *args);
-PyAPI_FUNC(_PyInitError) _PyPreConfig_Write(_PyPreConfig *config);
+PyAPI_FUNC(_PyInitError) _PyPreConfig_Write(const _PyPreConfig *config);
 
 
 /* --- _PyCoreConfig ---------------------------------------------- */

Include/internal/pycore_pymem.h

@@ -179,6 +179,15 @@ static inline int _PyMem_IsPtrFreed(void *ptr)
 #endif
 }
 
+PyAPI_FUNC(int) _PyMem_GetAllocatorName(
+    const char *name,
+    PyMemAllocatorName *allocator);
+
+/* Configure the Python memory allocators.
+   Pass PYMEM_ALLOCATOR_DEFAULT to use default allocators.
+   PYMEM_ALLOCATOR_NOT_SET does nothing. */
+PyAPI_FUNC(int) _PyMem_SetupAllocators(PyMemAllocatorName allocator);
+
 #ifdef __cplusplus
 }
 #endif

Lib/test/test_embed.py

@@ -13,6 +13,9 @@
 
 
 MS_WINDOWS = (os.name == 'nt')
+PYMEM_ALLOCATOR_NOT_SET = 0
+PYMEM_ALLOCATOR_DEBUG = 2
+PYMEM_ALLOCATOR_MALLOC = 3
 
 
 class EmbeddingTestsMixin:
@@ -272,7 +275,7 @@ class InitConfigTests(EmbeddingTestsMixin, unittest.TestCase):
     # Mark config which should be get by get_default_config()
     GET_DEFAULT_CONFIG = object()
     DEFAULT_PRE_CONFIG = {
-        'allocator': None,
+        'allocator': PYMEM_ALLOCATOR_NOT_SET,
         'coerce_c_locale': 0,
         'coerce_c_locale_warn': 0,
         'utf8_mode': 0,
@@ -564,7 +567,7 @@ def test_init_global_config(self):
 
     def test_init_from_config(self):
         preconfig = {
-            'allocator': 'malloc',
+            'allocator': PYMEM_ALLOCATOR_MALLOC,
             'utf8_mode': 1,
         }
         config = {
@@ -608,7 +611,7 @@ def test_init_from_config(self):
         self.check_config("init_from_config", config, preconfig)
 
     INIT_ENV_PRECONFIG = {
-        'allocator': 'malloc',
+        'allocator': PYMEM_ALLOCATOR_MALLOC,
     }
     INIT_ENV_CONFIG = {
         'use_hash_seed': 1,
@@ -633,23 +636,23 @@ def test_init_env(self):
 
     def test_init_env_dev_mode(self):
         preconfig = dict(self.INIT_ENV_PRECONFIG,
-                      allocator='debug')
+                      allocator=PYMEM_ALLOCATOR_DEBUG)
         config = dict(self.INIT_ENV_CONFIG,
                       dev_mode=1,
                       warnoptions=['default'])
         self.check_config("init_env_dev_mode", config, preconfig)
 
     def test_init_env_dev_mode_alloc(self):
         preconfig = dict(self.INIT_ENV_PRECONFIG,
-                         allocator='malloc')
+                         allocator=PYMEM_ALLOCATOR_MALLOC)
         config = dict(self.INIT_ENV_CONFIG,
                       dev_mode=1,
                       warnoptions=['default'])
         self.check_config("init_env_dev_mode_alloc", config, preconfig)
 
     def test_init_dev_mode(self):
         preconfig = {
-            'allocator': 'debug',
+            'allocator': PYMEM_ALLOCATOR_DEBUG,
         }
         config = {
             'faulthandler': 1,

Modules/_testcapimodule.c

@@ -4297,7 +4297,7 @@ pymem_malloc_without_gil(PyObject *self, PyObject *args)
 static PyObject*
 test_pymem_getallocatorsname(PyObject *self, PyObject *args)
 {
-    const char *name = _PyMem_GetAllocatorsName();
+    const char *name = _PyMem_GetCurrentAllocatorName();
     if (name == NULL) {
         PyErr_SetString(PyExc_RuntimeError, "cannot get allocators name");
         return NULL;

Objects/obmalloc.c

@@ -268,49 +268,94 @@ _PyMem_SetDefaultAllocator(PyMemAllocatorDomain domain,
 
 
 int
-_PyMem_SetupAllocators(const char *opt)
+_PyMem_GetAllocatorName(const char *name, PyMemAllocatorName *allocator)
 {
-    if (opt == NULL || *opt == '\0') {
+    if (name == NULL || *name == '\0') {
         /* PYTHONMALLOC is empty or is not set or ignored (-E/-I command line
-           options): use default memory allocators */
-        opt = "default";
+           nameions): use default memory allocators */
+        *allocator = PYMEM_ALLOCATOR_DEFAULT;
     }
+    else if (strcmp(name, "default") == 0) {
+        *allocator = PYMEM_ALLOCATOR_DEFAULT;
+    }
+    else if (strcmp(name, "debug") == 0) {
+        *allocator = PYMEM_ALLOCATOR_DEBUG;
+    }
+#ifdef WITH_PYMALLOC
+    else if (strcmp(name, "pymalloc") == 0) {
+        *allocator = PYMEM_ALLOCATOR_PYMALLOC;
+    }
+    else if (strcmp(name, "pymalloc_debug") == 0) {
+        *allocator = PYMEM_ALLOCATOR_PYMALLOC_DEBUG;
+    }
+#endif
+    else if (strcmp(name, "malloc") == 0) {
+        *allocator = PYMEM_ALLOCATOR_MALLOC;
+    }
+    else if (strcmp(name, "malloc_debug") == 0) {
+        *allocator = PYMEM_ALLOCATOR_MALLOC_DEBUG;
+    }
+    else {
+        /* unknown allocator */
+        return -1;
+    }
+    return 0;
+}
+
 
-    if (strcmp(opt, "default") == 0) {
+int
+_PyMem_SetupAllocators(PyMemAllocatorName allocator)
+{
+    switch (allocator) {
+    case PYMEM_ALLOCATOR_NOT_SET:
+        /* do nothing */
+        break;
+
+    case PYMEM_ALLOCATOR_DEFAULT:
         (void)_PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, NULL);
         (void)_PyMem_SetDefaultAllocator(PYMEM_DOMAIN_MEM, NULL);
         (void)_PyMem_SetDefaultAllocator(PYMEM_DOMAIN_OBJ, NULL);
-    }
-    else if (strcmp(opt, "debug") == 0) {
+        break;
+
+    case PYMEM_ALLOCATOR_DEBUG:
         (void)pymem_set_default_allocator(PYMEM_DOMAIN_RAW, 1, NULL);
         (void)pymem_set_default_allocator(PYMEM_DOMAIN_MEM, 1, NULL);
         (void)pymem_set_default_allocator(PYMEM_DOMAIN_OBJ, 1, NULL);
-    }
+        break;
+
 #ifdef WITH_PYMALLOC
-    else if (strcmp(opt, "pymalloc") == 0 || strcmp(opt, "pymalloc_debug") == 0) {
+    case PYMEM_ALLOCATOR_PYMALLOC:
+    case PYMEM_ALLOCATOR_PYMALLOC_DEBUG:
+    {
         PyMemAllocatorEx malloc_alloc = MALLOC_ALLOC;
         PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &malloc_alloc);
 
         PyMemAllocatorEx pymalloc = PYMALLOC_ALLOC;
         PyMem_SetAllocator(PYMEM_DOMAIN_MEM, &pymalloc);
         PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &pymalloc);
 
-        if (strcmp(opt, "pymalloc_debug") == 0) {
+        if (allocator == PYMEM_ALLOCATOR_PYMALLOC_DEBUG) {
             PyMem_SetupDebugHooks();
         }
+        break;
     }
 #endif
-    else if (strcmp(opt, "malloc") == 0 || strcmp(opt, "malloc_debug") == 0) {
+
+    case PYMEM_ALLOCATOR_MALLOC:
+    case PYMEM_ALLOCATOR_MALLOC_DEBUG:
+    {
         PyMemAllocatorEx malloc_alloc = MALLOC_ALLOC;
         PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &malloc_alloc);
         PyMem_SetAllocator(PYMEM_DOMAIN_MEM, &malloc_alloc);
         PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &malloc_alloc);
 
-        if (strcmp(opt, "malloc_debug") == 0) {
+        if (allocator == PYMEM_ALLOCATOR_MALLOC_DEBUG) {
             PyMem_SetupDebugHooks();
         }
+        break;
     }
-    else {
+
+    default:
         /* unknown allocator */
         return -1;
     }
@@ -326,7 +371,7 @@ pymemallocator_eq(PyMemAllocatorEx *a, PyMemAllocatorEx *b)
 
 
 const char*
-_PyMem_GetAllocatorsName(void)
+_PyMem_GetCurrentAllocatorName(void)
 {
     PyMemAllocatorEx malloc_alloc = MALLOC_ALLOC;
 #ifdef WITH_PYMALLOC

Programs/_testembed.c

@@ -379,7 +379,7 @@ static int test_init_from_config(void)
     _PyPreConfig preconfig = _PyPreConfig_INIT;
 
     putenv("PYTHONMALLOC=malloc_debug");
-    preconfig.allocator = "malloc";
+    preconfig.allocator = PYMEM_ALLOCATOR_MALLOC;
 
     putenv("PYTHONUTF8=0");
     Py_UTF8Mode = 0;

Python/coreconfig.c

@@ -2021,26 +2021,22 @@ core_read_precmdline(_PyCoreConfig *config, _PyPreCmdline *precmdline)
     _PyPreConfig preconfig = _PyPreConfig_INIT;
     if (_PyPreConfig_Copy(&preconfig, &_PyRuntime.preconfig) < 0) {
         err = _Py_INIT_NO_MEMORY();
-        goto done;
+        return err;
     }
 
     _PyCoreConfig_GetCoreConfig(&preconfig, config);
 
     err = _PyPreCmdline_Read(precmdline, &preconfig);
     if (_Py_INIT_FAILED(err)) {
-        goto done;
+        return err;
     }
 
     if (_PyPreCmdline_SetCoreConfig(precmdline, config) < 0) {
         err = _Py_INIT_NO_MEMORY();
-        goto done;
+        return err;
     }
 
-    err = _Py_INIT_OK();
-
-done:
-    _PyPreConfig_Clear(&preconfig);
-    return err;
+    return _Py_INIT_OK();
 }