diff --git a/Lib/pickle.py b/Lib/pickle.py index ab63cb21103..3d4c1a7c984 100644 --- a/Lib/pickle.py +++ b/Lib/pickle.py @@ -1093,11 +1093,16 @@ class _Pickler: (obj, module_name, name)) if self.proto >= 2: - code = _extension_registry.get((module_name, name)) - if code: - assert code > 0 + code = _extension_registry.get((module_name, name), _NoValue) + if code is not _NoValue: if code <= 0xff: - write(EXT1 + pack("extension_registry, - extension_key); + if (PyDict_GetItemRef(st->extension_registry, extension_key, &code_obj) < 0) { + Py_DECREF(extension_key); + goto error; + } Py_DECREF(extension_key); - /* The object is not registered in the extension registry. - This is the most likely code path. */ if (code_obj == NULL) { - if (PyErr_Occurred()) { - goto error; - } + /* The object is not registered in the extension registry. + This is the most likely code path. */ goto gen_global; } - /* XXX: pickle.py doesn't check neither the type, nor the range - of the value returned by the extension_registry. It should for - consistency. */ - - /* Verify code_obj has the right type and value. */ - if (!PyLong_Check(code_obj)) { - PyErr_Format(st->PicklingError, - "Can't pickle %R: extension code %R isn't an integer", - obj, code_obj); - goto error; - } - code = PyLong_AS_LONG(code_obj); + code = PyLong_AsLong(code_obj); + Py_DECREF(code_obj); if (code <= 0 || code > 0x7fffffffL) { + /* Should never happen in normal circumstances, since the type and + the value of the code are checked in copyreg.add_extension(). */ if (!PyErr_Occurred()) - PyErr_Format(st->PicklingError, "Can't pickle %R: extension " - "code %ld is out of range", obj, code); + PyErr_Format(PyExc_RuntimeError, "extension code %ld is out of range", code); goto error; }