Added __enter__ and __exit__ functions to HKEY object

Added ExpandEnvironmentStrings to the _winreg module.

1 files changed, 25 insertions, 20 deletions

diff --git a/Lib/test/test_winreg.py b/Lib/test/test_winreg.py
index 85a0aac..f2ec2a8 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)