Added __enter__ and __exit__ functions to HKEY object

Added ExpandEnvironmentStrings to the _winreg module.

4 files changed, 104 insertions, 20 deletions

diff --git a/Doc/library/_winreg.rst b/Doc/library/_winreg.rst
index 7a5aa38858..2544b5a60c 100644
--- a/Doc/library/_winreg.rst
+++ b/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
+
diff --git a/Lib/test/test_winreg.py b/Lib/test/test_winreg.py
index 85a0aacb95..f2ec2a8c0f 100644
--- a/Lib/test/test_winreg.py
+++ b/Lib/test/test_winreg.py
@@ -80,26 +80,26 @@ class WinregTests(unittest.TestCase):
 
         key = OpenKey(root_key, test_key_name)
         # Read the sub-keys
-        sub_key = OpenKey(key, "sub_key")
-        # Check I can enumerate over the values.
-        index = 0
-        while 1:
-            try:
-                data = EnumValue(sub_key, index)
-            except EnvironmentError:
-                break
-            self.assertEquals(data in test_data, True,
-                              "Didn't read back the correct test data")
-            index = index + 1
-        self.assertEquals(index, len(test_data),
-                          "Didn't read the correct number of items")
-        # Check I can directly access each item
-        for value_name, value_data, value_type in test_data:
-            read_val, read_typ = QueryValueEx(sub_key, value_name)
-            self.assertEquals(read_val, value_data,
-                              "Could not directly read the value")
-            self.assertEquals(read_typ, value_type,
-                              "Could not directly read the value")
+        with OpenKey(key, "sub_key") as sub_key:
+            # Check I can enumerate over the values.
+            index = 0
+            while 1:
+                try:
+                    data = EnumValue(sub_key, index)
+                except EnvironmentError:
+                    break
+                self.assertEquals(data in test_data, True,
+                                  "Didn't read back the correct test data")
+                index = index + 1
+            self.assertEquals(index, len(test_data),
+                              "Didn't read the correct number of items")
+            # Check I can directly access each item
+            for value_name, value_data, value_type in test_data:
+                read_val, read_typ = QueryValueEx(sub_key, value_name)
+                self.assertEquals(read_val, value_data,
+                                  "Could not directly read the value")
+                self.assertEquals(read_typ, value_type,
+                                  "Could not directly read the value")
         sub_key.Close()
         # Enumerate our main key.
         read_val = EnumKey(key, 0)
@@ -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)
 
diff --git a/Misc/NEWS b/Misc/NEWS
index e1745a7488..c77ce62cfe 100644
--- a/Misc/NEWS
+++ b/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.
 
diff --git a/PC/_winreg.c b/PC/_winreg.c
index 9ddbb236a9..19f928e6c1 100644
--- a/PC/_winreg.c
+++ b/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}
 };
 
@@ -1118,6 +1140,39 @@ PyEnumValue(PyObject *self, PyObject *args)
 }
 
 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)
 {
 	HKEY hKey;
@@ -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},