Index: Doc/library/_winreg.rst =================================================================== --- Doc/library/_winreg.rst (revision 78055) +++ Doc/library/_winreg.rst (working copy) @@ -108,6 +108,24 @@ raised, indicating, no more values are available. +.. function:: IterKey(key) + + Return an iterator over the subkeys of an open registry key. + + *key* is an already open key, or any one of the predefined :const:`HKEY_\*` + constants. + + The function is an iterable version of :func:`EnumKey`. + + >>> for subkey in IterKey(HKEY_CURRENT_USER): + ... print subkey + ... + AppEvents + Console + Control Panel + ... + + .. function:: EnumValue(key, index) Enumerates values of an open registry key, returning a tuple. @@ -135,8 +153,26 @@ | ``2`` | An integer that identifies the type of the | | | value data | +-------+--------------------------------------------+ + + +.. function:: IterValue(key) + Return an iterator over the values of an open registry key. The iterator + yields the same tuples returned by :func:`EnumValue`. + *key* is an already open key, or any one of the predefined :const:`HKEY_\*` + constants. + + The function is an iterable version of :func:`EnumValue`. + + >>> key = OpenKey(HKEY_CURRENT_USER, "Console") + >>> for value in IterValue(key): + ... print value + ... + ('ColorTable00', 0, 4) + ('ColorTable01', 8388608, 4) + + .. function:: ExpandEnvironmentStrings(unicode) Expands environment strings %NAME% in unicode string like const:`REG_EXPAND_SZ`:: Index: Lib/test/test_winreg.py =================================================================== --- Lib/test/test_winreg.py (revision 78055) +++ Lib/test/test_winreg.py (working copy) @@ -189,7 +189,30 @@ r = ExpandEnvironmentStrings(u"%windir%\\test") self.assertEqual(type(r), unicode) self.assertEqual(r, os.environ["windir"] + "\\test") + + def _test_iterable(self, orig, iterable): + # Test that an iterable version of the orig function yields the + # same values. + with OpenKey(HKEY_CURRENT_USER, "Software") as key: + idx = 0 + values = [] + while True: + try: + values.append(orig(key, idx)) + idx += 1 + except WindowsError: + break + itervalues = [] + for val in iterable(key): + itervalues.append(val) + + self.assertSameElements(values, itervalues) + + def testIterables(self): + self._test_iterable(EnumKey, IterKey) + self._test_iterable(EnumValue, IterValue) + def test_main(): test_support.run_unittest(WinregTests) Index: PC/_winreg.c =================================================================== --- PC/_winreg.c (revision 78055) +++ PC/_winreg.c (working copy) @@ -132,6 +132,15 @@ "It is typically called repeatedly until an EnvironmentError exception is\n" "raised, indicating no more values are available."); +PyDoc_STRVAR(IterKey_doc, +"iterable = IterKey(key) - Return an iterable sequence of subkeys.\n" +"\n" +"key is an already open key, or any one of the predefined HKEY_* constants.\n" +"\n" +"The function an allows the user to iterate over subkeys using\n" +"\"for subkey in IterKey(key)\", analagous to writing your own loop\n" +"over EnumKey."); + PyDoc_STRVAR(EnumValue_doc, "tuple = EnumValue(key, index) - Enumerates values of an open registry key.\n" "key is an already open key, or any one of the predefined HKEY_* constants.\n" @@ -147,6 +156,15 @@ " on the underlying registry type.\n" "data_type is an integer that identifies the type of the value data."); +PyDoc_STRVAR(IterValue_doc, +"iterable = IterValue(key) - Return an iterable sequence of values.\n" +"\n" +"key is an already open key, or any one of the predefined HKEY_* constants.\n" +"\n" +"The function an allows the user to iterate over values using\n" +"\"for value in IterValue(key)\", analagous to writing your own loop\n" +"over EnumValue."); + PyDoc_STRVAR(ExpandEnvironmentStrings_doc, "string = ExpandEnvironmentStrings(string) - Expand environment vars.\n"); @@ -1078,6 +1096,41 @@ return retStr; /* can be NULL */ } +typedef struct { + PyObject_HEAD + int index; + HKEY key; +} IterKey; + +PyObject * +PyIterKey_IterNext(PyObject *self) +{ + IterKey *iter = (IterKey*)self; + + long rc; + PyObject *retStr; + char tmpbuf[256]; /* max key name length is 255 */ + DWORD len = sizeof(tmpbuf); /* includes NULL terminator */ + + Py_BEGIN_ALLOW_THREADS + rc = RegEnumKeyEx(iter->key, iter->index, tmpbuf, &len, + NULL, NULL, NULL, NULL); + Py_END_ALLOW_THREADS + + if (rc == ERROR_NO_MORE_ITEMS) + { + Py_DECREF(iter); + PyErr_SetNone(PyExc_StopIteration); + return NULL; + } + else if (rc != ERROR_SUCCESS) + return PyErr_SetFromWindowsErrWithFunction(rc, "RegEnumKeyEx"); + + iter->index++; + retStr = PyString_FromStringAndSize(tmpbuf, len); + return retStr; /* can be NULL */ +} + static PyObject * PyEnumValue(PyObject *self, PyObject *args) { @@ -1144,6 +1197,71 @@ return retVal; } +PyObject * +PyIterValue_IterNext(PyObject *self) +{ + IterKey *iter = (IterKey*)self; + + long rc; + char *retValueBuf; + char *retDataBuf; + DWORD retValueSize; + DWORD retDataSize; + DWORD typ; + PyObject *obData; + PyObject *retVal; + + if ((rc = RegQueryInfoKey(iter->key, NULL, NULL, NULL, NULL, NULL, + NULL, NULL, &retValueSize, &retDataSize, + NULL, NULL)) + != ERROR_SUCCESS) + return PyErr_SetFromWindowsErrWithFunction(rc, + "RegQueryInfoKey"); + ++retValueSize; /* include null terminators */ + ++retDataSize; + retValueBuf = (char *)PyMem_Malloc(retValueSize); + if (retValueBuf == NULL) + return PyErr_NoMemory(); + retDataBuf = (char *)PyMem_Malloc(retDataSize); + if (retDataBuf == NULL) { + PyMem_Free(retValueBuf); + return PyErr_NoMemory(); + } + + Py_BEGIN_ALLOW_THREADS + rc = RegEnumValue(iter->key, iter->index, retValueBuf, &retValueSize, + NULL, &typ, (BYTE *)retDataBuf, &retDataSize); + Py_END_ALLOW_THREADS + + if (rc == ERROR_NO_MORE_ITEMS) + { + Py_DECREF(iter); + PyMem_Free(retValueBuf); + PyMem_Free(retDataBuf); + PyErr_SetNone(PyExc_StopIteration); + return NULL; + } + else if (rc != ERROR_SUCCESS) { + retVal = PyErr_SetFromWindowsErrWithFunction(rc, + "PyRegEnumValue"); + goto fail; + } + + obData = Reg2Py(retDataBuf, retDataSize, typ); + if (obData == NULL) { + retVal = NULL; + goto fail; + } + retVal = Py_BuildValue("sOi", retValueBuf, obData, typ); + Py_DECREF(obData); + iter->index++; + + fail: + PyMem_Free(retValueBuf); + PyMem_Free(retDataBuf); + return retVal; +} + static PyObject * PyExpandEnvironmentStrings(PyObject *self, PyObject *args) { @@ -1572,6 +1690,114 @@ return PyBool_FromLong(rc); } +static PyTypeObject IterKeyType = { + PyObject_HEAD_INIT(NULL) + 0, /*ob_size*/ + "_winreg._IterKeyType", /*tp_name*/ + sizeof(IterKey), /*tp_basicsize*/ + 0, /*tp_itemsize*/ + 0, /*tp_dealloc*/ + 0, /*tp_print*/ + 0, /*tp_getattr*/ + 0, /*tp_setattr*/ + 0, /*tp_compare*/ + 0, /*tp_repr*/ + 0, /*tp_as_number*/ + 0, /*tp_as_sequence*/ + 0, /*tp_as_mapping*/ + 0, /*tp_hash */ + 0, /*tp_call*/ + 0, /*tp_str*/ + 0, /*tp_getattro*/ + 0, /*tp_setattro*/ + 0, /*tp_as_buffer*/ + Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_ITER, /*tp_flags*/ + "RegEnumKeyEx iterator", /*tp_doc*/ + 0, /*tp_traverse*/ + 0, /*tp_clear*/ + 0, /*tp_richcompare*/ + 0, /*tp_weaklistoffset*/ + PyObject_SelfIter, /*tp_iter*/ + (iternextfunc)PyIterKey_IterNext /*tp_iternext*/ +}; + +static PyTypeObject IterValueType = { + PyObject_HEAD_INIT(NULL) + 0, /*ob_size*/ + "_winreg._IterValueType", /*tp_name*/ + sizeof(IterKey), /*tp_basicsize*/ + 0, /*tp_itemsize*/ + 0, /*tp_dealloc*/ + 0, /*tp_print*/ + 0, /*tp_getattr*/ + 0, /*tp_setattr*/ + 0, /*tp_compare*/ + 0, /*tp_repr*/ + 0, /*tp_as_number*/ + 0, /*tp_as_sequence*/ + 0, /*tp_as_mapping*/ + 0, /*tp_hash */ + 0, /*tp_call*/ + 0, /*tp_str*/ + 0, /*tp_getattro*/ + 0, /*tp_setattro*/ + 0, /*tp_as_buffer*/ + Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_ITER, /*tp_flags*/ + "RegEnumValueEx iterator", /*tp_doc*/ + 0, /*tp_traverse*/ + 0, /*tp_clear*/ + 0, /*tp_richcompare*/ + 0, /*tp_weaklistoffset*/ + PyObject_SelfIter, /*tp_iter*/ + (iternextfunc)PyIterValue_IterNext /*tp_iternext*/ +}; + +static PyObject * +PyIterKey(PyObject *self, PyObject *args) +{ + HKEY hKey; + PyObject *obKey; + IterKey *ik; + + if (!PyArg_ParseTuple(args, "O:IterKey", &obKey)) + return NULL; + + if (!PyHKEY_AsHKEY(obKey, &hKey, FALSE)) + return NULL; + + ik = PyObject_New(IterKey, &IterKeyType); + if (!ik) + return NULL; + + Py_INCREF(ik); + ik->index = 0; + ik->key = hKey; + return (PyObject*)ik; +} + +static PyObject * +PyIterValue(PyObject *self, PyObject *args) +{ + HKEY hKey; + PyObject *obKey; + IterKey *ik; + + if (!PyArg_ParseTuple(args, "O:IterValue", &obKey)) + return NULL; + + if (!PyHKEY_AsHKEY(obKey, &hKey, FALSE)) + return NULL; + + ik = PyObject_New(IterKey, &IterValueType); + if (!ik) + return NULL; + + Py_INCREF(ik); + ik->index = 0; + ik->key = hKey; + return (PyObject*)ik; +} + static struct PyMethodDef winreg_methods[] = { {"CloseKey", PyCloseKey, METH_VARARGS, CloseKey_doc}, {"ConnectRegistry", PyConnectRegistry, METH_VARARGS, ConnectRegistry_doc}, @@ -1581,7 +1807,9 @@ {"DisableReflectionKey", PyDisableReflectionKey, METH_VARARGS, DisableReflectionKey_doc}, {"EnableReflectionKey", PyEnableReflectionKey, METH_VARARGS, EnableReflectionKey_doc}, {"EnumKey", PyEnumKey, METH_VARARGS, EnumKey_doc}, + {"IterKey", PyIterKey, METH_VARARGS, IterKey_doc}, {"EnumValue", PyEnumValue, METH_VARARGS, EnumValue_doc}, + {"IterValue", PyIterValue, METH_VARARGS, IterValue_doc}, {"ExpandEnvironmentStrings", PyExpandEnvironmentStrings, METH_VARARGS, ExpandEnvironmentStrings_doc }, {"FlushKey", PyFlushKey, METH_VARARGS, FlushKey_doc}, @@ -1623,6 +1851,15 @@ PyMODINIT_FUNC init_winreg(void) { PyObject *m, *d; + + IterKeyType.tp_new = PyType_GenericNew; + if (PyType_Ready(&IterKeyType) < 0) + return; + + IterValueType.tp_new = PyType_GenericNew; + if (PyType_Ready(&IterValueType) < 0) + return; + m = Py_InitModule3("_winreg", winreg_methods, module_doc); if (m == NULL) return; @@ -1639,6 +1876,14 @@ PyExc_WindowsError) != 0) return; + Py_INCREF(&IterKeyType); + if (PyDict_SetItemString(d, "IterKeyType", + (PyObject *)&IterKeyType) != 0) + + Py_INCREF(&IterValueType); + if (PyDict_SetItemString(d, "IterValueType", + (PyObject *)&IterValueType) != 0) + /* Add the relevant constants */ ADD_KEY(HKEY_CLASSES_ROOT); ADD_KEY(HKEY_CURRENT_USER);