Added __enter__ and __exit__ functions to HKEY object

Added ExpandEnvironmentStrings to the _winreg module.
This commit is contained in:
Christian Heimes 2008-01-08 15:46:10 +00:00
parent 41f278ffa5
commit b39a756afd
4 changed files with 104 additions and 20 deletions

View File

@ -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) .. function:: FlushKey(key)
Writes all the attributes of a key to the registry. 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 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. 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

View File

@ -80,26 +80,26 @@ def ReadTestData(self, root_key):
key = OpenKey(root_key, test_key_name) key = OpenKey(root_key, test_key_name)
# Read the sub-keys # Read the sub-keys
sub_key = OpenKey(key, "sub_key") with OpenKey(key, "sub_key") as sub_key:
# Check I can enumerate over the values. # Check I can enumerate over the values.
index = 0 index = 0
while 1: while 1:
try: try:
data = EnumValue(sub_key, index) data = EnumValue(sub_key, index)
except EnvironmentError: except EnvironmentError:
break break
self.assertEquals(data in test_data, True, self.assertEquals(data in test_data, True,
"Didn't read back the correct test data") "Didn't read back the correct test data")
index = index + 1 index = index + 1
self.assertEquals(index, len(test_data), self.assertEquals(index, len(test_data),
"Didn't read the correct number of items") "Didn't read the correct number of items")
# Check I can directly access each item # Check I can directly access each item
for value_name, value_data, value_type in test_data: for value_name, value_data, value_type in test_data:
read_val, read_typ = QueryValueEx(sub_key, value_name) read_val, read_typ = QueryValueEx(sub_key, value_name)
self.assertEquals(read_val, value_data, self.assertEquals(read_val, value_data,
"Could not directly read the value") "Could not directly read the value")
self.assertEquals(read_typ, value_type, self.assertEquals(read_typ, value_type,
"Could not directly read the value") "Could not directly read the value")
sub_key.Close() sub_key.Close()
# Enumerate our main key. # Enumerate our main key.
read_val = EnumKey(key, 0) read_val = EnumKey(key, 0)
@ -161,6 +161,11 @@ def testRemoteMachineRegistryWorks(self):
remote_key = ConnectRegistry(self.remote_name, HKEY_CURRENT_USER) remote_key = ConnectRegistry(self.remote_name, HKEY_CURRENT_USER)
self.TestAll(remote_key) 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(): def test_main():
test_support.run_unittest(WinregTests) test_support.run_unittest(WinregTests)

View File

@ -944,6 +944,10 @@ Library
Extension Modules 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 - Issue #1646: Make socket support TIPC. The socket module now has support
for TIPC under Linux, see http://tipc.sf.net/ for more information. for TIPC under Linux, see http://tipc.sf.net/ for more information.

View File

@ -47,6 +47,7 @@ PyDoc_STRVAR(module_doc,
"DeleteValue() - Removes a named value from the specified registry key.\n" "DeleteValue() - Removes a named value from the specified registry key.\n"
"EnumKey() - Enumerates subkeys of the specified open registry key.\n" "EnumKey() - Enumerates subkeys of the specified open registry key.\n"
"EnumValue() - Enumerates values 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" "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" "LoadKey() - Creates a subkey under HKEY_USER or HKEY_LOCAL_MACHINE and stores\n"
" registration information from a specified file into that subkey.\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" " on the underlying registry type.\n"
"data_type is an integer that identifies the type of the value data."); "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, PyDoc_STRVAR(FlushKey_doc,
"FlushKey(key) - Writes all the attributes of a key to the registry.\n" "FlushKey(key) - Writes all the attributes of a key to the registry.\n"
"\n" "\n"
@ -518,9 +522,27 @@ PyHKEY_DetachMethod(PyObject *self, PyObject *args)
return PyLong_FromVoidPtr(ret); 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[] = { static struct PyMethodDef PyHKEY_methods[] = {
{"Close", PyHKEY_CloseMethod, METH_VARARGS, PyHKEY_Close_doc}, {"Close", PyHKEY_CloseMethod, METH_VARARGS, PyHKEY_Close_doc},
{"Detach", PyHKEY_DetachMethod, METH_VARARGS, PyHKEY_Detach_doc}, {"Detach", PyHKEY_DetachMethod, METH_VARARGS, PyHKEY_Detach_doc},
{"__enter__", (PyCFunction)PyHKEY_Enter, METH_NOARGS, NULL},
{"__exit__", PyHKEY_Exit, METH_VARARGS, NULL},
{NULL} {NULL}
}; };
@ -1117,6 +1139,39 @@ PyEnumValue(PyObject *self, PyObject *args)
return retVal; 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 * static PyObject *
PyFlushKey(PyObject *self, PyObject *args) PyFlushKey(PyObject *self, PyObject *args)
{ {
@ -1412,6 +1467,8 @@ static struct PyMethodDef winreg_methods[] = {
{"DeleteValue", PyDeleteValue, METH_VARARGS, DeleteValue_doc}, {"DeleteValue", PyDeleteValue, METH_VARARGS, DeleteValue_doc},
{"EnumKey", PyEnumKey, METH_VARARGS, EnumKey_doc}, {"EnumKey", PyEnumKey, METH_VARARGS, EnumKey_doc},
{"EnumValue", PyEnumValue, METH_VARARGS, EnumValue_doc}, {"EnumValue", PyEnumValue, METH_VARARGS, EnumValue_doc},
{"ExpandEnvironmentStrings", PyExpandEnvironmentStrings, METH_VARARGS,
ExpandEnvironmentStrings_doc },
{"FlushKey", PyFlushKey, METH_VARARGS, FlushKey_doc}, {"FlushKey", PyFlushKey, METH_VARARGS, FlushKey_doc},
{"LoadKey", PyLoadKey, METH_VARARGS, LoadKey_doc}, {"LoadKey", PyLoadKey, METH_VARARGS, LoadKey_doc},
{"OpenKey", PyOpenKey, METH_VARARGS, OpenKey_doc}, {"OpenKey", PyOpenKey, METH_VARARGS, OpenKey_doc},