diff --git a/python/extension_dict.c b/python/extension_dict.c index c2bd4d6572..0111fa60d2 100644 --- a/python/extension_dict.c +++ b/python/extension_dict.c @@ -128,6 +128,8 @@ static int PyUpb_ExtensionDict_AssignSubscript(PyObject* _self, PyObject* key, } } +static PyObject* PyUpb_ExtensionIterator_New(PyObject* _ext_dict); + static PyMethodDef PyUpb_ExtensionDict_Methods[] = { {"_FindExtensionByName", PyUpb_ExtensionDict_FindExtensionByName, METH_O, "Finds an extension by name."}, @@ -142,7 +144,7 @@ static PyType_Slot PyUpb_ExtensionDict_Slots[] = { //{Py_tp_getset, PyUpb_ExtensionDict_Getters}, //{Py_tp_hash, PyObject_HashNotImplemented}, //{Py_tp_richcompare, PyUpb_ExtensionDict_RichCompare}, - //{Py_tp_iter, PyUpb_ExtensionDict_GetIter}, + {Py_tp_iter, PyUpb_ExtensionIterator_New}, {Py_sq_contains, PyUpb_ExtensionDict_Contains}, {Py_sq_length, PyUpb_ExtensionDict_Length}, {Py_mp_length, PyUpb_ExtensionDict_Length}, @@ -164,23 +166,55 @@ static PyType_Spec PyUpb_ExtensionDict_Spec = { typedef struct { PyObject_HEAD - PyObject* arena; + PyObject* msg; + size_t iter; } PyUpb_ExtensionIterator; -static void PyUpb_ExtensionIterator_Dealloc(PyUpb_ExtensionDict* self) { - PyUpb_Dealloc(self); +static PyObject* PyUpb_ExtensionIterator_New(PyObject* _ext_dict) { + PyUpb_ExtensionDict* ext_dict = (PyUpb_ExtensionDict*)_ext_dict; + PyUpb_ModuleState* state = PyUpb_ModuleState_Get(); + PyUpb_ExtensionIterator* iter = + (void*)PyType_GenericAlloc(state->extension_iterator_type, 0); + iter->msg = ext_dict->msg; + iter->iter = UPB_MSG_BEGIN; + Py_INCREF(iter->msg); + return &iter->ob_base; +} + +static void PyUpb_ExtensionIterator_Dealloc(void* _self) { + PyUpb_ExtensionIterator* self = (PyUpb_ExtensionIterator*)_self; + Py_DECREF(&self->msg); + PyUpb_Dealloc(_self); +} + +PyObject* PyUpb_ExtensionIterator_IterNext(PyObject* _self) { + PyUpb_ExtensionIterator* self = (PyUpb_ExtensionIterator*)_self; + /* + if (self->version != self->map->version) { + return PyErr_Format(PyExc_RuntimeError, "Map modified during iteration."); + } + */ + upb_msg* msg = PyUpb_CMessage_GetIfReified(self->msg); + if (!msg) return NULL; + const upb_msgdef* m = PyUpb_CMessage_GetMsgdef(self->msg); + const upb_symtab* symtab = upb_filedef_symtab(upb_msgdef_file(m)); + while (true) { + const upb_fielddef* f; + upb_msgval val; + if (!upb_msg_next(msg, m, symtab, &f, &val, &self->iter)) return NULL; + if (upb_fielddef_isextension(f)) return PyUpb_FieldDescriptor_Get(f); + } } static PyType_Slot PyUpb_ExtensionIterator_Slots[] = { {Py_tp_dealloc, PyUpb_ExtensionIterator_Dealloc}, {Py_tp_iter, PyObject_SelfIter}, - //{Py_tp_iter, PyUpb_ExtensionIterator_GetIter}, - //{Py_tp_iternext, PyUpb_ExtensionIterator_IterNext}, + {Py_tp_iternext, PyUpb_ExtensionIterator_IterNext}, {0, NULL}}; static PyType_Spec PyUpb_ExtensionIterator_Spec = { PYUPB_MODULE_NAME ".ExtensionIterator", // tp_name - sizeof(PyUpb_ExtensionDict), // tp_basicsize + sizeof(PyUpb_ExtensionIterator), // tp_basicsize 0, // tp_itemsize Py_TPFLAGS_DEFAULT, // tp_flags PyUpb_ExtensionIterator_Slots, diff --git a/python/pb_unit_tests/reflection_test_wrapper.py b/python/pb_unit_tests/reflection_test_wrapper.py index 83ba3dab07..7969abaeaa 100644 --- a/python/pb_unit_tests/reflection_test_wrapper.py +++ b/python/pb_unit_tests/reflection_test_wrapper.py @@ -26,7 +26,10 @@ from google.protobuf.internal import reflection_test import unittest +# The tests depend on a specific iteration order for extensions, which is not +# specified or guaranteed. reflection_test.Proto2ReflectionTest.testExtensionIter.__unittest_expecting_failure__ = True + reflection_test.Proto2ReflectionTest.testIsInitialized.__unittest_expecting_failure__ = True reflection_test.Proto2ReflectionTest.testListFieldsAndExtensions.__unittest_expecting_failure__ = True reflection_test.Proto2ReflectionTest.testRepeatedCompositeReverse_Empty.__unittest_expecting_failure__ = True