Merge pull request RustPython#1793 from RustPython/coolreader18/winreg-module

youknowone · web-flow · commit a6c84c299403 · 2020-03-04T22:59:00.000+09:00
Add the winreg stdlib module
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Lib/importlib/_bootstrap_external.py b/Lib/importlib/_bootstrap_external.py
@@ -1605,8 +1605,7 @@ def _setup(_bootstrap_module):
     setattr(self_module, '_weakref', weakref_module)
 
     # Directly load the winreg module (needed during bootstrap).
-    # XXX RustPython TODO: winreg module
-    if builtin_os == 'nt' and False:
+    if builtin_os == 'nt':
         winreg_module = _bootstrap._builtin_from_name('winreg')
         setattr(self_module, '_winreg', winreg_module)
 
diff --git a/vm/Cargo.toml b/vm/Cargo.toml
@@ -94,6 +94,9 @@ rustyline = "6.0"
 [target.'cfg(not(any(target_arch = "wasm32", target_os = "redox")))'.dependencies]
 dns-lookup = "1.0"
 
+[target.'cfg(windows)'.dependencies]
+winreg = "0.7"
+
 [target."cfg(windows)".dependencies.winapi]
 version = "0.3"
 features = ["winsock2", "handleapi", "ws2def", "std", "winbase"]
diff --git a/vm/src/stdlib/mod.rs b/vm/src/stdlib/mod.rs
@@ -52,6 +52,8 @@ pub mod signal;
 mod subprocess;
 #[cfg(windows)]
 mod winapi;
+#[cfg(windows)]
+mod winreg;
 #[cfg(not(target_arch = "wasm32"))]
 mod zlib;
 
@@ -135,6 +137,7 @@ pub fn get_module_inits() -> HashMap<String, StdlibInitFunc> {
     #[cfg(windows)]
     {
         modules.insert("_winapi".to_owned(), Box::new(winapi::make_module));
+        modules.insert("winreg".to_owned(), Box::new(winreg::make_module));
     }
 
     modules
diff --git a/vm/src/stdlib/winreg.rs b/vm/src/stdlib/winreg.rs
@@ -0,0 +1,312 @@
+#![allow(non_snake_case)]
+
+use std::cell::{Ref, RefCell};
+use std::convert::TryInto;
+use std::io;
+
+use super::os;
+use crate::function::OptionalArg;
+use crate::obj::objstr::PyStringRef;
+use crate::obj::objtype::PyClassRef;
+use crate::pyobject::{PyClassImpl, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject};
+use crate::VirtualMachine;
+
+use winapi::shared::winerror;
+use winreg::{enums::RegType, RegKey, RegValue};
+
+#[pyclass]
+#[derive(Debug)]
+struct PyHKEY {
+    key: RefCell<RegKey>,
+}
+type PyHKEYRef = PyRef<PyHKEY>;
+
+impl PyValue for PyHKEY {
+    fn class(vm: &VirtualMachine) -> PyClassRef {
+        vm.class("winreg", "HKEYType")
+    }
+}
+
+#[pyimpl]
+impl PyHKEY {
+    fn new(key: RegKey) -> Self {
+        Self {
+            key: RefCell::new(key),
+        }
+    }
+
+    fn key(&self) -> Ref<RegKey> {
+        self.key.borrow()
+    }
+
+    #[pymethod]
+    fn Close(&self) {
+        let null_key = RegKey::predef(0 as winreg::HKEY);
+        let key = self.key.replace(null_key);
+        drop(key);
+    }
+    #[pymethod]
+    fn Detach(&self) -> usize {
+        let null_key = RegKey::predef(0 as winreg::HKEY);
+        let key = self.key.replace(null_key);
+        let handle = key.raw_handle();
+        std::mem::forget(key);
+        handle as usize
+    }
+
+    #[pymethod(magic)]
+    fn bool(&self) -> bool {
+        !self.key().raw_handle().is_null()
+    }
+    #[pymethod(magic)]
+    fn enter(zelf: PyRef<Self>) -> PyRef<Self> {
+        zelf
+    }
+    #[pymethod(magic)]
+    fn exit(&self, _cls: PyObjectRef, _exc: PyObjectRef, _tb: PyObjectRef) {
+        self.Close();
+    }
+}
+
+enum Hkey {
+    PyHKEY(PyHKEYRef),
+    Constant(winreg::HKEY),
+}
+impl TryFromObject for Hkey {
+    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
+        obj.downcast()
+            .map(Self::PyHKEY)
+            .or_else(|o| usize::try_from_object(vm, o).map(|i| Self::Constant(i as winreg::HKEY)))
+    }
+}
+impl Hkey {
+    fn with_key<R>(&self, f: impl FnOnce(&RegKey) -> R) -> R {
+        match self {
+            Self::PyHKEY(py) => f(&py.key()),
+            Self::Constant(hkey) => {
+                let k = RegKey::predef(*hkey);
+                let res = f(&k);
+                std::mem::forget(k);
+                res
+            }
+        }
+    }
+}
+
+fn winreg_OpenKey(
+    key: Hkey,
+    subkey: Option<PyStringRef>,
+    reserved: OptionalArg<i32>,
+    access: OptionalArg<u32>,
+    vm: &VirtualMachine,
+) -> PyResult<PyHKEY> {
+    let reserved = reserved.unwrap_or(0);
+    let access = access.unwrap_or(winreg::enums::KEY_READ);
+    if reserved != 0 {
+        // RegKey::open_subkey* doesn't have a reserved param, so this'll do
+        return Err(vm.new_value_error("reserved param must be 0".to_owned()));
+    }
+
+    let subkey = subkey.as_ref().map_or("", |s| s.as_str());
+    let key = key
+        .with_key(|k| k.open_subkey_with_flags(subkey, access))
+        .map_err(|e| os::convert_io_error(vm, e))?;
+
+    Ok(PyHKEY::new(key))
+}
+
+fn winreg_QueryValue(
+    key: Hkey,
+    subkey: Option<PyStringRef>,
+    vm: &VirtualMachine,
+) -> PyResult<String> {
+    let subkey = subkey.as_ref().map_or("", |s| s.as_str());
+    key.with_key(|k| k.get_value(subkey))
+        .map_err(|e| os::convert_io_error(vm, e))
+}
+
+fn winreg_QueryValueEx(
+    key: Hkey,
+    subkey: Option<PyStringRef>,
+    vm: &VirtualMachine,
+) -> PyResult<(PyObjectRef, usize)> {
+    let subkey = subkey.as_ref().map_or("", |s| s.as_str());
+    key.with_key(|k| k.get_raw_value(subkey))
+        .map_err(|e| os::convert_io_error(vm, e))
+        .and_then(|regval| {
+            let ty = regval.vtype.clone() as usize;
+            Ok((reg_to_py(regval, vm)?, ty))
+        })
+}
+
+fn winreg_EnumKey(key: Hkey, index: u32, vm: &VirtualMachine) -> PyResult<String> {
+    key.with_key(|k| k.enum_keys().nth(index as usize))
+        .unwrap_or_else(|| {
+            Err(io::Error::from_raw_os_error(
+                winerror::ERROR_NO_MORE_ITEMS as i32,
+            ))
+        })
+        .map_err(|e| os::convert_io_error(vm, e))
+}
+
+fn winreg_EnumValue(
+    key: Hkey,
+    index: u32,
+    vm: &VirtualMachine,
+) -> PyResult<(String, PyObjectRef, usize)> {
+    key.with_key(|k| k.enum_values().nth(index as usize))
+        .unwrap_or_else(|| {
+            Err(io::Error::from_raw_os_error(
+                winerror::ERROR_NO_MORE_ITEMS as i32,
+            ))
+        })
+        .map_err(|e| os::convert_io_error(vm, e))
+        .and_then(|(name, value)| {
+            let ty = value.vtype.clone() as usize;
+            Ok((name, reg_to_py(value, vm)?, ty))
+        })
+}
+
+fn winreg_CloseKey(key: Hkey) {
+    match key {
+        Hkey::PyHKEY(py) => py.Close(),
+        Hkey::Constant(hkey) => drop(RegKey::predef(hkey)),
+    }
+}
+
+fn reg_to_py(value: RegValue, vm: &VirtualMachine) -> PyResult {
+    macro_rules! bytes_to_int {
+        ($int:ident, $f:ident, $name:ident) => {{
+            let i = if value.bytes.is_empty() {
+                Ok(0 as $int)
+            } else {
+                (&*value.bytes).try_into().map($int::$f).map_err(|_| {
+                    vm.new_value_error(format!("{} value is wrong length", stringify!(name)))
+                })
+            };
+            i.map(|i| vm.new_int(i))
+        }};
+    };
+    let bytes_to_wide = |b: &[u8]| -> Option<&[u16]> {
+        if b.len() % 2 == 0 {
+            Some(unsafe { std::slice::from_raw_parts(b.as_ptr().cast(), b.len() / 2) })
+        } else {
+            None
+        }
+    };
+    match value.vtype {
+        RegType::REG_DWORD => bytes_to_int!(u32, from_ne_bytes, REG_DWORD),
+        RegType::REG_DWORD_BIG_ENDIAN => bytes_to_int!(u32, from_be_bytes, REG_DWORD_BIG_ENDIAN),
+        RegType::REG_QWORD => bytes_to_int!(u64, from_ne_bytes, REG_DWORD),
+        // RegType::REG_QWORD_BIG_ENDIAN => bytes_to_int!(u64, from_be_bytes, REG_DWORD_BIG_ENDIAN),
+        RegType::REG_SZ | RegType::REG_EXPAND_SZ => {
+            let wide_slice = bytes_to_wide(&value.bytes).ok_or_else(|| {
+                vm.new_value_error("REG_SZ string doesn't have an even byte length".to_owned())
+            })?;
+            let nul_pos = wide_slice
+                .iter()
+                .position(|w| *w == 0)
+                .unwrap_or_else(|| wide_slice.len());
+            let s = String::from_utf16_lossy(&wide_slice[..nul_pos]);
+            Ok(vm.new_str(s))
+        }
+        RegType::REG_MULTI_SZ => {
+            if value.bytes.is_empty() {
+                return Ok(vm.ctx.new_list(vec![]));
+            }
+            let wide_slice = bytes_to_wide(&value.bytes).ok_or_else(|| {
+                vm.new_value_error(
+                    "REG_MULTI_SZ string doesn't have an even byte length".to_owned(),
+                )
+            })?;
+            let wide_slice = if let Some((0, rest)) = wide_slice.split_last() {
+                rest
+            } else {
+                wide_slice
+            };
+            let strings = wide_slice
+                .split(|c| *c == 0)
+                .map(|s| vm.new_str(String::from_utf16_lossy(s)))
+                .collect();
+            Ok(vm.ctx.new_list(strings))
+        }
+        RegType::REG_BINARY | _ => {
+            if value.bytes.is_empty() {
+                Ok(vm.get_none())
+            } else {
+                Ok(vm.ctx.new_bytes(value.bytes))
+            }
+        }
+    }
+}
+
+pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
+    let ctx = &vm.ctx;
+    let hkey_type = PyHKEY::make_class(ctx);
+    let module = py_module!(vm, "winreg", {
+        "HKEYType" => hkey_type,
+        "OpenKey" => ctx.new_function(winreg_OpenKey),
+        "QueryValue" => ctx.new_function(winreg_QueryValue),
+        "QueryValueEx" => ctx.new_function(winreg_QueryValueEx),
+        "EnumKey" => ctx.new_function(winreg_EnumKey),
+        "EnumValue" => ctx.new_function(winreg_EnumValue),
+        "CloseKey" => ctx.new_function(winreg_CloseKey),
+    });
+
+    macro_rules! add_constants {
+        (hkey, $($name:ident),*$(,)?) => {
+            extend_module!(vm, module, {
+                $((stringify!($name)) => ctx.new_int(winreg::enums::$name as usize)),*
+            })
+        };
+        (winnt, $($name:ident),*$(,)?) => {
+            extend_module!(vm, module, {
+                $((stringify!($name)) => ctx.new_int(winapi::um::winnt::$name)),*
+            })
+        };
+    }
+
+    add_constants!(
+        hkey,
+        HKEY_CLASSES_ROOT,
+        HKEY_CURRENT_USER,
+        HKEY_LOCAL_MACHINE,
+        HKEY_USERS,
+        HKEY_PERFORMANCE_DATA,
+        HKEY_CURRENT_CONFIG,
+        HKEY_DYN_DATA,
+    );
+    add_constants!(
+        winnt,
+        // access rights
+        KEY_ALL_ACCESS,
+        KEY_WRITE,
+        KEY_READ,
+        KEY_EXECUTE,
+        KEY_QUERY_VALUE,
+        KEY_SET_VALUE,
+        KEY_CREATE_SUB_KEY,
+        KEY_ENUMERATE_SUB_KEYS,
+        KEY_NOTIFY,
+        KEY_CREATE_LINK,
+        KEY_WOW64_64KEY,
+        KEY_WOW64_32KEY,
+        // value types
+        REG_BINARY,
+        REG_DWORD,
+        REG_DWORD_LITTLE_ENDIAN,
+        REG_DWORD_BIG_ENDIAN,
+        REG_EXPAND_SZ,
+        REG_LINK,
+        REG_MULTI_SZ,
+        REG_NONE,
+        REG_QWORD,
+        REG_QWORD_LITTLE_ENDIAN,
+        REG_RESOURCE_LIST,
+        REG_FULL_RESOURCE_DESCRIPTOR,
+        REG_RESOURCE_REQUIREMENTS_LIST,
+        REG_SZ,
+    );
+
+    module
+}

Original file line number	Diff line number	Diff line change
`@@ -52,6 +52,8 @@ pub mod signal;`
`52`	`52`	`mod subprocess;`
`53`	`53`	`#[cfg(windows)]`
`54`	`54`	`mod winapi;`
	`55`	`+#[cfg(windows)]`
	`56`	`+mod winreg;`
`55`	`57`	`#[cfg(not(target_arch = "wasm32"))]`
`56`	`58`	`mod zlib;`
`57`	`59`
`@@ -135,6 +137,7 @@ pub fn get_module_inits() -> HashMap<String, StdlibInitFunc> {`
`135`	`137`	`#[cfg(windows)]`
`136`	`138`	`{`
`137`	`139`	`modules.insert("_winapi".to_owned(), Box::new(winapi::make_module));`
	`140`	`+ modules.insert("winreg".to_owned(), Box::new(winreg::make_module));`
`138`	`141`	`}`
`139`	`142`
`140`	`143`	`modules`