1 files changed, 48 insertions, 26 deletions

diff --git a/Python/initconfig.c b/Python/initconfig.c
index 1733ed05a3e..8fc36b1c0bc 100644
--- a/Python/initconfig.c
+++ b/Python/initconfig.c
@@ -7,7 +7,7 @@
 #include "pycore_pathconfig.h"    // _Py_path_config
 #include "pycore_pyerrors.h"      // _PyErr_GetRaisedException()
 #include "pycore_pylifecycle.h"   // _Py_PreInitializeFromConfig()
-#include "pycore_pymem.h"         // _PyMem_SetDefaultAllocator()
+#include "pycore_pymem.h"         // _PyMem_DefaultRawMalloc()
 #include "pycore_pystate.h"       // _PyThreadState_GET()
 #include "pycore_pystats.h"       // _Py_StatsOn()
 #include "pycore_sysmodule.h"     // _PySys_SetIntMaxStrDigits()
@@ -619,53 +619,87 @@ _PyWideStringList_CheckConsistency(const PyWideStringList *list)
 #endif   /* Py_DEBUG */
 
 
-void
-_PyWideStringList_Clear(PyWideStringList *list)
+static void
+_PyWideStringList_ClearEx(PyWideStringList *list,
+                          bool use_default_allocator)
 {
     assert(_PyWideStringList_CheckConsistency(list));
     for (Py_ssize_t i=0; i < list->length; i++) {
-        PyMem_RawFree(list->items[i]);
+        if (use_default_allocator) {
+            _PyMem_DefaultRawFree(list->items[i]);
+        }
+        else {
+            PyMem_RawFree(list->items[i]);
+        }
+    }
+    if (use_default_allocator) {
+        _PyMem_DefaultRawFree(list->items);
+    }
+    else {
+        PyMem_RawFree(list->items);
     }
-    PyMem_RawFree(list->items);
     list->length = 0;
     list->items = NULL;
 }
 
+void
+_PyWideStringList_Clear(PyWideStringList *list)
+{
+    _PyWideStringList_ClearEx(list, false);
+}
 
-int
-_PyWideStringList_Copy(PyWideStringList *list, const PyWideStringList *list2)
+static int
+_PyWideStringList_CopyEx(PyWideStringList *list,
+                         const PyWideStringList *list2,
+                         bool use_default_allocator)
 {
     assert(_PyWideStringList_CheckConsistency(list));
     assert(_PyWideStringList_CheckConsistency(list2));
 
     if (list2->length == 0) {
-        _PyWideStringList_Clear(list);
+        _PyWideStringList_ClearEx(list, use_default_allocator);
         return 0;
     }
 
     PyWideStringList copy = _PyWideStringList_INIT;
 
     size_t size = list2->length * sizeof(list2->items[0]);
-    copy.items = PyMem_RawMalloc(size);
+    if (use_default_allocator) {
+        copy.items = _PyMem_DefaultRawMalloc(size);
+    }
+    else {
+        copy.items = PyMem_RawMalloc(size);
+    }
     if (copy.items == NULL) {
         return -1;
     }
 
     for (Py_ssize_t i=0; i < list2->length; i++) {
-        wchar_t *item = _PyMem_RawWcsdup(list2->items[i]);
+        wchar_t *item;
+        if (use_default_allocator) {
+            item = _PyMem_DefaultRawWcsdup(list2->items[i]);
+        }
+        else {
+            item = _PyMem_RawWcsdup(list2->items[i]);
+        }
         if (item == NULL) {
-            _PyWideStringList_Clear(&copy);
+            _PyWideStringList_ClearEx(&copy, use_default_allocator);
             return -1;
         }
         copy.items[i] = item;
         copy.length = i + 1;
     }
 
-    _PyWideStringList_Clear(list);
+    _PyWideStringList_ClearEx(list, use_default_allocator);
     *list = copy;
     return 0;
 }
 
+int
+_PyWideStringList_Copy(PyWideStringList *list, const PyWideStringList *list2)
+{
+    return _PyWideStringList_CopyEx(list, list2, false);
+}
 
 PyStatus
 PyWideStringList_Insert(PyWideStringList *list,
@@ -789,12 +823,7 @@ _PyWideStringList_AsTuple(const PyWideStringList *list)
 void
 _Py_ClearArgcArgv(void)
 {
-    PyMemAllocatorEx old_alloc;
-    _PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-
-    _PyWideStringList_Clear(&_PyRuntime.orig_argv);
-
-    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
+    _PyWideStringList_ClearEx(&_PyRuntime.orig_argv, true);
 }
 
 
@@ -802,17 +831,10 @@ static int
 _Py_SetArgcArgv(Py_ssize_t argc, wchar_t * const *argv)
 {
     const PyWideStringList argv_list = {.length = argc, .items = (wchar_t **)argv};
-    int res;
-
-    PyMemAllocatorEx old_alloc;
-    _PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
 
     // XXX _PyRuntime.orig_argv only gets cleared by Py_Main(),
     // so it currently leaks for embedders.
-    res = _PyWideStringList_Copy(&_PyRuntime.orig_argv, &argv_list);
-
-    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
-    return res;
+    return _PyWideStringList_CopyEx(&_PyRuntime.orig_argv, &argv_list, true);
 }