1 files changed, 46 insertions, 39 deletions

diff --git a/PC/_subprocess.c b/PC/_subprocess.c
index 2338f30..93e51d3 100644
--- a/PC/_subprocess.c
+++ b/PC/_subprocess.c
@@ -329,12 +329,9 @@ gethandle(PyObject* obj, char* name)
 static PyObject*
 getenvironment(PyObject* environment)
 {
-    int i;
-    Py_ssize_t envsize;
-    PyObject* out = NULL;
-    PyObject* keys;
-    PyObject* values;
-    Py_UNICODE* p;
+    Py_ssize_t i, envsize, totalsize;
+    Py_UCS4 *buffer = NULL, *p, *end;
+    PyObject *keys, *values, *res;
 
     /* convert environment dictionary to windows enviroment string */
     if (! PyMapping_Check(environment)) {
@@ -350,14 +347,8 @@ getenvironment(PyObject* environment)
     if (!keys || !values)
         goto error;
 
-    out = PyUnicode_FromUnicode(NULL, 2048);
-    if (! out)
-        goto error;
-
-    p = PyUnicode_AS_UNICODE(out);
-
+    totalsize = 1; /* trailing null character */
     for (i = 0; i < envsize; i++) {
-        Py_ssize_t ksize, vsize, totalsize;
         PyObject* key = PyList_GET_ITEM(keys, i);
         PyObject* value = PyList_GET_ITEM(values, i);
 
@@ -366,36 +357,42 @@ getenvironment(PyObject* environment)
                 "environment can only contain strings");
             goto error;
         }
-        ksize = PyUnicode_GET_SIZE(key);
-        vsize = PyUnicode_GET_SIZE(value);
-        totalsize = (p - PyUnicode_AS_UNICODE(out)) + ksize + 1 +
-                                                     vsize + 1 + 1;
-        if (totalsize > PyUnicode_GET_SIZE(out)) {
-            Py_ssize_t offset = p - PyUnicode_AS_UNICODE(out);
-            PyUnicode_Resize(&out, totalsize + 1024);
-            p = PyUnicode_AS_UNICODE(out) + offset;
-        }
-        Py_UNICODE_COPY(p, PyUnicode_AS_UNICODE(key), ksize);
-        p += ksize;
+        totalsize += PyUnicode_GET_LENGTH(key) + 1;    /* +1 for '=' */
+        totalsize += PyUnicode_GET_LENGTH(value) + 1;  /* +1 for '\0' */
+    }
+
+    buffer = PyMem_Malloc(totalsize * sizeof(Py_UCS4));
+    if (! buffer)
+        goto error;
+    p = buffer;
+    end = buffer + totalsize;
+
+    for (i = 0; i < envsize; i++) {
+        PyObject* key = PyList_GET_ITEM(keys, i);
+        PyObject* value = PyList_GET_ITEM(values, i);
+        if (!PyUnicode_AsUCS4(key, p, end - p, 0))
+            goto error;
+        p += PyUnicode_GET_LENGTH(key);
         *p++ = '=';
-        Py_UNICODE_COPY(p, PyUnicode_AS_UNICODE(value), vsize);
-        p += vsize;
+        if (!PyUnicode_AsUCS4(value, p, end - p, 0))
+            goto error;
+        p += PyUnicode_GET_LENGTH(value);
         *p++ = '\0';
     }
 
     /* add trailing null byte */
     *p++ = '\0';
-    PyUnicode_Resize(&out, p - PyUnicode_AS_UNICODE(out));
-
-    /* PyObject_Print(out, stdout, 0); */
+    assert(p == end);
 
     Py_XDECREF(keys);
     Py_XDECREF(values);
 
-    return out;
+    res = PyUnicode_FromKindAndData(PyUnicode_4BYTE_KIND, buffer, p - buffer);
+    PyMem_Free(buffer);
+    return res;
 
  error:
-    Py_XDECREF(out);
+    PyMem_Free(buffer);
     Py_XDECREF(keys);
     Py_XDECREF(values);
     return NULL;
@@ -420,15 +417,16 @@ sp_CreateProcess(PyObject* self, PyObject* args)
     PROCESS_INFORMATION pi;
     STARTUPINFOW si;
     PyObject* environment;
+    wchar_t *wenvironment;
 
-    Py_UNICODE* application_name;
-    Py_UNICODE* command_line;
+    wchar_t* application_name;
+    wchar_t* command_line;
     PyObject* process_attributes; /* ignored */
     PyObject* thread_attributes; /* ignored */
     int inherit_handles;
     int creation_flags;
     PyObject* env_mapping;
-    Py_UNICODE* current_directory;
+    wchar_t* current_directory;
     PyObject* startup_info;
 
     if (! PyArg_ParseTuple(args, "ZZOOiiOZO:CreateProcess",
@@ -456,12 +454,20 @@ sp_CreateProcess(PyObject* self, PyObject* args)
     if (PyErr_Occurred())
         return NULL;
 
-    if (env_mapping == Py_None)
-        environment = NULL;
-    else {
+    if (env_mapping != Py_None) {
         environment = getenvironment(env_mapping);
         if (! environment)
             return NULL;
+        wenvironment = PyUnicode_AsUnicode(environment);
+        if (wenvironment == NULL)
+        {
+            Py_XDECREF(environment);
+            return NULL;
+        }
+    }
+    else {
+        environment = NULL;
+        wenvironment = NULL;
     }
 
     Py_BEGIN_ALLOW_THREADS
@@ -471,7 +477,7 @@ sp_CreateProcess(PyObject* self, PyObject* args)
                            NULL,
                            inherit_handles,
                            creation_flags | CREATE_UNICODE_ENVIRONMENT,
-                           environment ? PyUnicode_AS_UNICODE(environment) : NULL,
+                           wenvironment,
                            current_directory,
                            &si,
                            &pi);
@@ -609,7 +615,7 @@ sp_GetModuleFileName(PyObject* self, PyObject* args)
     if (! result)
         return PyErr_SetFromWindowsErr(GetLastError());
 
-    return PyUnicode_FromUnicode(filename, Py_UNICODE_strlen(filename));
+    return PyUnicode_FromWideChar(filename, wcslen(filename));
 }
 
 static PyMethodDef sp_functions[] = {
@@ -682,6 +688,7 @@ PyInit__subprocess()
     defint(d, "SW_HIDE", SW_HIDE);
     defint(d, "INFINITE", INFINITE);
     defint(d, "WAIT_OBJECT_0", WAIT_OBJECT_0);
+    defint(d, "WAIT_TIMEOUT", WAIT_TIMEOUT);
     defint(d, "CREATE_NEW_CONSOLE", CREATE_NEW_CONSOLE);
     defint(d, "CREATE_NEW_PROCESS_GROUP", CREATE_NEW_PROCESS_GROUP);