[pywin32-checkins] /hgroot/pywin32/pywin32: Add GetEnvironmentVariableW and SetEnvi...

[<pyw...@li...>]

changeset 31014eee4df2 in /hgroot/pywin32/pywin32
details: http://pywin32.hg.sourceforge.net/hgweb/pywin32/pywin32/hgroot/pywin32/pywin32?cmd=changeset;node=31014eee4df2
summary: Add GetEnvironmentVariableW and SetEnvironmentVariableW

diffstat:

 win32/src/win32apimodule.cpp |  101 ++++++++++++++++++++++++++++++++++++++++--
 1 files changed, 95 insertions(+), 6 deletions(-)

diffs (164 lines):

diff -r e2cdf93c66b5 -r 31014eee4df2 win32/src/win32apimodule.cpp
--- a/win32/src/win32apimodule.cpp	Sat Jan 25 14:43:34 2014 +1100
+++ b/win32/src/win32apimodule.cpp	Fri Feb 14 17:29:03 2014 -0500
@@ -388,21 +388,22 @@
 	return Py_None;
 }
 
-// @pymethod string|win32api|GetEnvironmentVariable|Retrieves the value of an environment variable.
-// @comm Returns None if environment variable is not found
+// @pymethod str|win32api|GetEnvironmentVariable|Retrieves the value of an environment variable.
+// @rdesc Returns None if environment variable is not found
 static PyObject *
 PyGetEnvironmentVariable( PyObject *self, PyObject *args )
 {
 	TCHAR *szVar;
 	PyObject *obVar, *ret=NULL;
 	if (!PyArg_ParseTuple(args, "O:GetEnvironmentVariable", 
-	           &obVar)) // @pyparm string|variable||The variable to get
+	           &obVar)) // @pyparm str|variable||The variable to get
 		return NULL;
 	if (!PyWinObject_AsTCHAR(obVar, &szVar, FALSE))
 		return NULL;
 	// @pyseeapi GetEnvironmentVariable
 	PyW32_BEGIN_ALLOW_THREADS
 	DWORD size = GetEnvironmentVariable(szVar, NULL, 0);
+	PyW32_END_ALLOW_THREADS
 	TCHAR *pResult = NULL;
 	if (!size){
 		Py_INCREF(Py_None);
@@ -411,13 +412,14 @@
 	else{
 		pResult = (TCHAR *)malloc(sizeof(TCHAR) * size);
 		if (pResult==NULL)
-			PyErr_NoMemory();	// ??? does this need to hold thread lock ???
+			PyErr_NoMemory();
 		else{
+			PyW32_BEGIN_ALLOW_THREADS
 			GetEnvironmentVariable(szVar, pResult, size);
+			PyW32_END_ALLOW_THREADS
 			ret = PyWinObject_FromTCHAR(pResult);
 			}
 		}
-	PyW32_END_ALLOW_THREADS
 
 	PyWinObject_FreeTCHAR(szVar);
 	if (pResult)
@@ -425,6 +427,60 @@
 	return ret;
 }
 
+// @pymethod <o PyUnicode>|win32api|GetEnvironmentVariableW|Retrieves the unicode value of an environment variable.
+// @rdesc Returns None if environment variable is not found
+// @pyseeapi GetEnvironmentVariableW
+static PyObject *
+PyGetEnvironmentVariableW( PyObject *self, PyObject *args )
+{
+	TmpWCHAR Name;
+	PyObject *obName;
+	if (!PyArg_ParseTuple(args, "O:GetEnvironmentVariableW", 
+	    &obName)) // @pyparm str|Name||The variable to retrieve
+		return NULL;
+	if (!PyWinObject_AsWCHAR(obName, &Name))
+		return NULL;
+
+	DWORD returned_size, allocated_size = 0;
+	WCHAR *pResult = NULL;
+	PyObject *ret = NULL;
+	// Call in loop to account for race condition where env var is changed between calls
+	while(TRUE){
+		if (pResult)
+			free(pResult);
+		if (allocated_size){
+			// returned_size includes NULL terminator
+			pResult = (WCHAR *)malloc(allocated_size * sizeof(WCHAR));
+			if (pResult == NULL){
+				PyErr_NoMemory();
+				break;
+				}
+			}
+		Py_BEGIN_ALLOW_THREADS
+		returned_size = GetEnvironmentVariableW(Name, pResult, allocated_size);
+		Py_END_ALLOW_THREADS
+		if (!returned_size){
+			DWORD err = GetLastError();
+			if (err == ERROR_ENVVAR_NOT_FOUND){
+				Py_INCREF(Py_None);
+				ret = Py_None;
+				}
+			else
+				PyWin_SetAPIError("GetEnvironmentVariableW", err);
+			break;
+			}
+		// Var may have been changed between calls, check that value still fits in buffer
+		if (returned_size < allocated_size){
+			ret = PyWinObject_FromWCHAR(pResult, returned_size);
+			break;
+			}
+		allocated_size = returned_size;
+		}
+	if (pResult)
+		free(pResult);
+	return ret;
+}
+
 // @pymethod |win32api|SetEnvironmentVariable|Creates, deletes, or changes the value of an environment variable.
 static PyObject *
 PySetEnvironmentVariable(PyObject *self, PyObject *args)
@@ -450,6 +506,27 @@
 	return ret;
 }
 
+// @pymethod |win32api|SetEnvironmentVariableW|Creates, deletes, or changes the value of an environment variable.
+static PyObject *
+PySetEnvironmentVariableW(PyObject *self, PyObject *args)
+{
+	TmpWCHAR Name, Value;
+	PyObject *obName, *obValue;
+	if (!PyArg_ParseTuple(args, "OO:SetEnvironmentVariableW", 
+		&obName,		// @pyparm str|Name||Name of the environment variable
+		&obValue))		// @pyparm str|Value||Value to be set, or None to remove variable
+		return NULL;
+	// @pyseeapi SetEnvironmentVariable
+	if (!PyWinObject_AsWCHAR(obName, &Name))
+		return NULL;
+	if (!PyWinObject_AsWCHAR(obValue, &Value, TRUE))
+		return NULL;
+	if (!SetEnvironmentVariableW(Name, Value))
+		return PyWin_SetAPIError("SetEnvironmentVariableW");
+	Py_INCREF(Py_None);
+	return Py_None;
+}
+
 // @pymethod string|win32api|ExpandEnvironmentStrings|Expands environment-variable strings and replaces them with their defined values. 
 static PyObject *
 PyExpandEnvironmentStrings( PyObject *self, PyObject *args )
@@ -6333,7 +6410,13 @@
 	{"GetDiskFreeSpaceEx",	PyGetDiskFreeSpaceEx, 1}, // @pymeth GetDiskFreeSpaceEx|Retrieves information about a disk.
 	{"GetDllDirectory",		PyGetDllDirectory, METH_NOARGS}, // @pymeth GetDllDirectory|Retrieves the DLL search path
 	{"GetDomainName",		PyGetDomainName, 1}, 	// @pymeth GetDomainName|Returns the current domain name
-	{"GetEnvironmentVariable", PyGetEnvironmentVariable, 1}, // @pymeth GetEnvironmentVariable|Retrieves the value of an environment variable.
+	// @pymeth GetEnvironmentVariable|Retrieves the value of an environment variable.
+#ifdef UNICODE
+	{"GetEnvironmentVariable", PyGetEnvironmentVariableW, 1}, 
+#else
+	{"GetEnvironmentVariable", PyGetEnvironmentVariable, 1},
+#endif
+	{"GetEnvironmentVariableW", PyGetEnvironmentVariableW, 1}, // @pymeth GetEnvironmentVariableW|Retrieves the value of an environment variable.
 	{"GetFileAttributes",   PyGetFileAttributes,1}, // @pymeth GetFileAttributes|Retrieves the attributes for the named file.
 	{"GetFileVersionInfo",	PyGetFileVersionInfo, 1}, //@pymeth GetFileVersionInfo|Retrieves string version info
 	{"GetFocus",            PyGetFocus,         1}, // @pymeth GetFocus|Retrieves the handle of the keyboard focus window associated with the thread that called the method. 
@@ -6472,7 +6555,13 @@
 	{"SetClassWord",       PySetClassWord,1}, // @pymeth SetClassWord|Replaces the specified 32-bit (long) value at the specified offset into the extra class memory for the window.
 	{"SetClassWord",       PySetWindowWord,1}, // @pymeth SetWindowWord|
 	{"SetCursor",           PySetCursor,1}, // @pymeth SetCursor|Set the cursor to the HCURSOR object.
+	// @pymeth SetEnvironmentVariable|Creates, deletes, or changes the value of an environment variable.
+#ifdef UNICODE
+	{"SetEnvironmentVariable", PySetEnvironmentVariableW,1},
+#else
 	{"SetEnvironmentVariable", PySetEnvironmentVariable,1}, // @pymeth SetEnvironmentVariable|Creates, deletes, or changes the value of an environment variable.
+#endif
+	{"SetEnvironmentVariableW", PySetEnvironmentVariableW,1}, // @pymeth SetEnvironmentVariableW|Creates, deletes, or changes the value of an environment variable.
 	{"SetHandleInformation",	PySetHandleInformation,1}, // @pymeth SetHandleInformation|Sets a handles's flags
 	{"SetStdHandle",	PySetStdHandle,	1}, // @pymeth SetStdHandle|Sets a handle for the standard input, standard output, or standard error device
 	{"SetSystemPowerState",	PySetSystemPowerState,	1}, // @pymeth SetSystemPowerState|Powers machine down to a suspended state