Added __enter__ and __exit__ functions to HKEY object

Added ExpandEnvironmentStrings to the _winreg module.

diff --git a/Doc/library/_winreg.rst b/Doc/library/_winreg.rst
index 7a5aa38858e1ed1aafb383d0dab1c163a2bf8330..2544b5a60c6ff1ec3d2df49243dc2b80f82f7731 100644 (file)
--- 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 85a0aacb95d8594c5136a18681c36bb56d616214..f2ec2a8c0f89c76517b3b3fefaf5f96776fe732d 100644 (file)
--- 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 e1745a74882cc74be1e09e81a09c4675bd676838..c77ce62cfefc3116041f9646e457edf1c4d16f8e 100644 (file)
--- 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 9ddbb236a99fa41175643199dc13456d45d78ea4..19f928e6c1dec5aeda1dafe246a39765761ba79d 100644 (file)
--- 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}
 };
 
@@ -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},