backport e070d661: implement winreg reflection on win64

2 files changed, 70 insertions, 10 deletions

diff --git a/pypy/module/_winreg/interp_winreg.py b/pypy/module/_winreg/interp_winreg.py
index c7af44d6da..226ab24a88 100644
--- a/pypy/module/_winreg/interp_winreg.py
+++ b/pypy/module/_winreg/interp_winreg.py
@@ -713,24 +713,45 @@ def DisableReflectionKey(space, w_key):
     a 32-bit Operating System.
     If the key is not on the reflection list, the function succeeds but has no effect.
     Disabling reflection for a key does not affect reflection of any subkeys."""
-    raise oefmt(space.w_NotImplementedError,
-                "not implemented on this platform")
+    if not rwinreg.reflection_supported:
+        raise oefmt(space.w_NotImplementedError,
+                    "not implemented on this platform")
+    else:
+        hkey = hkey_w(w_key, space)
+        ret = rwinreg.RegDisableReflectionKey(hkey)
+        if ret != 0:
+            raiseWindowsError(space, ret, 'RegDisableReflectionKey')
 
 def EnableReflectionKey(space, w_key):
     """Restores registry reflection for the specified disabled key.
     Will generally raise NotImplemented if executed on a 32-bit Operating System.
     Restoring reflection for a key does not affect reflection of any subkeys."""
-    raise oefmt(space.w_NotImplementedError,
-                "not implemented on this platform")
+    if not rwinreg.reflection_supported:
+        raise oefmt(space.w_NotImplementedError,
+                    "not implemented on this platform")
+    else:
+        hkey = hkey_w(w_key, space)
+        ret = rwinreg.RegEnableReflectionKey(hkey)
+        if ret != 0:
+            raiseWindowsError(space, ret, 'RegEnableReflectionKey')
 
 def QueryReflectionKey(space, w_key):
     """bool = QueryReflectionKey(hkey) - Determines the reflection state for the specified key.
     Will generally raise NotImplemented if executed on a 32-bit Operating System."""
-    raise oefmt(space.w_NotImplementedError,
-                "not implemented on this platform")
+    if not rwinreg.reflection_supported:
+        raise oefmt(space.w_NotImplementedError,
+                    "not implemented on this platform")
+    else:
+        hkey = hkey_w(w_key, space)
+        with lltype.scoped_alloc(rwin32.LPBOOL.TO, 1) as isDisabled:
+            ret = rwinreg.RegQueryReflectionKey(hkey, isDisabled)
+            if ret != 0:
+                raiseWindowsError(space, ret, 'RegQueryReflectionKey')
+            return space.newbool(intmask(isDisabled[0]) != 0)
 
-@unwrap_spec(subkey="text")
-def DeleteKeyEx(space, w_key, subkey):
+
+@unwrap_spec(sub_key="unicode", access=r_uint, reserved=int)
+def DeleteKeyEx(space, w_key, sub_key, access=rwinreg.KEY_WOW64_64KEY, reserved=0):
     """DeleteKeyEx(key, sub_key, sam, res) - Deletes the specified key.
 
     key is an already open key, or any one of the predefined HKEY_* constants.
@@ -744,5 +765,13 @@ def DeleteKeyEx(space, w_key, subkey):
     If the method succeeds, the entire key, including all of its values,
     is removed.  If the method fails, a WindowsError exception is raised.
     On unsupported Windows versions, NotImplementedError is raised."""
-    raise oefmt(space.w_NotImplementedError,
-                "not implemented on this platform")
+    if not rwinreg.reflection_supported:
+        raise oefmt(space.w_NotImplementedError,
+                    "not implemented on this platform")
+    else:
+        hkey = hkey_w(w_key, space)
+        with rffi.scoped_unicode2wcharp(sub_key) as wide_subkey:
+            c_subkey = rffi.cast(rffi.CWCHARP, wide_subkey)
+            ret = rwinreg.RegDeleteKeyExW(hkey, c_subkey, access, reserved)
+            if ret != 0:
+                raiseWindowsError(space, ret, 'RegDeleteKeyEx')
diff --git a/pypy/module/_winreg/test/test_winreg.py b/pypy/module/_winreg/test/test_winreg.py
index ae2d8ee227..73dad7bd15 100644
--- a/pypy/module/_winreg/test/test_winreg.py
+++ b/pypy/module/_winreg/test/test_winreg.py
@@ -255,3 +255,34 @@ class AppTestFfi:
             raises(NotImplementedError, DeleteKeyEx, self.root_key,
                    self.test_key_name)
 
+    def test_reflection(self):
+        import sys
+        from _winreg import DisableReflectionKey, EnableReflectionKey, \
+                           QueryReflectionKey, OpenKey, HKEY_LOCAL_MACHINE
+        # Adapted from lib-python test
+        if sys.getwindowsversion() < (5, 2) or sys.maxsize < 2**32:
+            skip("Requires 64-bit Windows")
+        # Test that we can call the query, enable, and disable functions
+        # on a key which isn't on the reflection list with no consequences.
+        with OpenKey(HKEY_LOCAL_MACHINE, "Software") as key:
+            # HKLM\Software is redirected but not reflected in all OSes
+            assert QueryReflectionKey(key)
+            assert EnableReflectionKey(key) is None
+            assert DisableReflectionKey(key) is None
+            assert QueryReflectionKey(key)
+
+    def test_named_arguments(self):
+        import sys
+        from _winreg import KEY_ALL_ACCESS, CreateKeyEx, DeleteKey, DeleteKeyEx, OpenKeyEx
+        with CreateKeyEx(key=self.root_key, sub_key=self.test_key_name,
+                         reserved=0, access=KEY_ALL_ACCESS) as ckey:
+            assert ckey.handle != 0
+        with OpenKeyEx(key=self.root_key, sub_key=self.test_key_name,
+                       reserved=0, access=KEY_ALL_ACCESS) as okey:
+            assert okey.handle != 0
+        if sys.getwindowsversion() >= (5, 2) and sys.maxsize > 2**32:
+            DeleteKeyEx(key=self.root_key, sub_key=self.test_key_name,
+                        access=KEY_ALL_ACCESS, reserved=0)
+        else:
+            DeleteKey(self.root_key, self.test_key_name)
+