diff --git a/Lib/pickle.py b/Lib/pickle.py index c4d6e658216..01c1a102794 100644 --- a/Lib/pickle.py +++ b/Lib/pickle.py @@ -397,6 +397,8 @@ def decode_long(data): return int.from_bytes(data, byteorder='little', signed=True) +_NoValue = object() + # Pickling machinery class _Pickler: @@ -1091,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); 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); + Py_INCREF(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; }