Added __enter__ and __exit__ functions to HKEY object

Added ExpandEnvironmentStrings to the _winreg module.

Doc/library/_winreg.rst

@@ -133,6 +133,16 @@ This module offers the following functions:
    +-------+--------------------------------------------+
 
 
+.. function:: ExpandEnvironmentStrings(unicode)
+
+   Expands environment strings %NAME% in unicode string like const:`REG_EXPAND_SZ`::
+   
+   >>> ExpandEnvironmentStrings(u"%windir%")
+   u"C:\\Windows"
+   
+   .. versionadded: 2.6
+
+
 .. function:: FlushKey(key)
 
    Writes all the attributes of a key to the registry.
@@ -418,3 +428,11 @@ handle, and also disconnect the Windows handle from the handle object.
    handle is not closed.  You would call this function when  you need the
    underlying Win32 handle to exist beyond the lifetime  of the handle object.
 
+.. method:: PyHKEY.__enter__()
+.. method:: PyHKEY.__exit__(*exc_info)
+   
+   The HKEY object implements __enter__ and __exit__ and thus supports the 
+   context protocol for the with statement.
+   
+   .. versionadded: 2.6
+

Lib/test/test_winreg.py

@@ -80,7 +80,7 @@ class WinregTests(unittest.TestCase):
 
         key = OpenKey(root_key, test_key_name)
         # Read the sub-keys
-        sub_key = OpenKey(key, "sub_key")
+        with OpenKey(key, "sub_key") as sub_key:
             # Check I can enumerate over the values.
             index = 0
             while 1:
@@ -161,6 +161,11 @@ class WinregTests(unittest.TestCase):
         remote_key = ConnectRegistry(self.remote_name, HKEY_CURRENT_USER)
         self.TestAll(remote_key)
 
+    def testExpandEnvironmentStrings(self):
+        r = ExpandEnvironmentStrings(u"%windir%\\test")
+        self.assertEqual(type(r), unicode)
+        self.assertEqual(r, os.environ["windir"] + "\\test")
+
 def test_main():
     test_support.run_unittest(WinregTests)
 

Misc/NEWS

@@ -944,6 +944,10 @@ Library
 Extension Modules
 -----------------
 
+- _winreg's HKEY object have gained __enter__ and __exit__ functions to support he
+  context manager protocol. The _winreg module also got a new function
+  ``ExpandEnvironmentStrings`` to expand REG_EXPAND_SZ keys.
+
 - Issue #1646: Make socket support TIPC. The socket module now has support
   for TIPC under Linux, see http://tipc.sf.net/ for more information.
 

PC/_winreg.c

@@ -47,6 +47,7 @@ PyDoc_STRVAR(module_doc,
 "DeleteValue() - Removes a named value from the specified registry key.\n"
 "EnumKey() - Enumerates subkeys of the specified open registry key.\n"
 "EnumValue() - Enumerates values of the specified open registry key.\n"
+"ExpandEnvironmentStrings() - Expand the env strings in a REG_EXPAND_SZ string.\n"
 "FlushKey() - Writes all the attributes of the specified key to the registry.\n"
 "LoadKey() - Creates a subkey under HKEY_USER or HKEY_LOCAL_MACHINE and stores\n"
 "            registration information from a specified file into that subkey.\n"
@@ -146,6 +147,9 @@ PyDoc_STRVAR(EnumValue_doc,
 " on the underlying registry type.\n"
 "data_type is an integer that identifies the type of the value data.");
 
+PyDoc_STRVAR(ExpandEnvironmentStrings_doc,
+"string = ExpandEnvironmentStrings(string) - Expand environment vars.\n");
+
 PyDoc_STRVAR(FlushKey_doc,
 "FlushKey(key) - Writes all the attributes of a key to the registry.\n"
 "\n"
@@ -518,9 +522,27 @@ PyHKEY_DetachMethod(PyObject *self, PyObject *args)
 	return PyLong_FromVoidPtr(ret);
 }
 
+static PyObject *
+PyHKEY_Enter(PyObject *self)
+{
+	Py_XINCREF(self);
+	return self;
+}
+
+static PyObject *
+PyHKEY_Exit(PyObject *self, PyObject *args)
+{
+	if (!PyHKEY_Close(self))
+		return NULL;
+	Py_RETURN_NONE;
+}
+
+
 static struct PyMethodDef PyHKEY_methods[] = {
 	{"Close",  PyHKEY_CloseMethod, METH_VARARGS, PyHKEY_Close_doc},
 	{"Detach", PyHKEY_DetachMethod, METH_VARARGS, PyHKEY_Detach_doc},
+	{"__enter__", (PyCFunction)PyHKEY_Enter, METH_NOARGS, NULL},
+	{"__exit__", PyHKEY_Exit, METH_VARARGS, NULL},
 	{NULL}
 };
 
@@ -1117,6 +1139,39 @@ PyEnumValue(PyObject *self, PyObject *args)
 	return retVal;
 }
 
+static PyObject *
+PyExpandEnvironmentStrings(PyObject *self, PyObject *args)
+{
+	Py_UNICODE *retValue = NULL;
+	Py_UNICODE *src;
+	DWORD retValueSize;
+	DWORD rc;
+	PyObject *o;
+
+	if (!PyArg_ParseTuple(args, "u:ExpandEnvironmentStrings", &src))
+		return NULL;
+
+	retValueSize = ExpandEnvironmentStringsW(src, retValue, 0);
+	if (retValueSize == 0) {
+		return PyErr_SetFromWindowsErrWithFunction(retValueSize,
+						"ExpandEnvironmentStrings");
+	}
+	retValue = (Py_UNICODE *)PyMem_Malloc(retValueSize * sizeof(Py_UNICODE));
+	if (retValue == NULL) {
+		return PyErr_NoMemory();
+	}
+
+	rc = ExpandEnvironmentStringsW(src, retValue, retValueSize);
+	if (rc == 0) {
+		PyMem_Free(retValue);
+		return PyErr_SetFromWindowsErrWithFunction(retValueSize,
+						"ExpandEnvironmentStrings");
+	}
+	o = PyUnicode_FromUnicode(retValue, wcslen(retValue));
+	PyMem_Free(retValue);
+	return o;
+}
+
 static PyObject *
 PyFlushKey(PyObject *self, PyObject *args)
 {
@@ -1412,6 +1467,8 @@ static struct PyMethodDef winreg_methods[] = {
 	{"DeleteValue",      PyDeleteValue,     METH_VARARGS, DeleteValue_doc},
 	{"EnumKey",          PyEnumKey,         METH_VARARGS, EnumKey_doc},
 	{"EnumValue",        PyEnumValue,       METH_VARARGS, EnumValue_doc},
+	{"ExpandEnvironmentStrings", PyExpandEnvironmentStrings, METH_VARARGS,
+		ExpandEnvironmentStrings_doc },
 	{"FlushKey",         PyFlushKey,        METH_VARARGS, FlushKey_doc},
 	{"LoadKey",          PyLoadKey,         METH_VARARGS, LoadKey_doc},
 	{"OpenKey",          PyOpenKey,         METH_VARARGS, OpenKey_doc},