[3.12] pythongh-122213: Add notes for pickle errors

When pickle encounters a serialization error deep in a nested data
structure, the resulting traceback gives no indication of where in
the object graph the failure occurred. This makes debugging difficult
for large or deeply nested objects.

Backport of c0c2aa7 from main. Each save() call site now catches
BaseException and attaches a note identifying the context, such as
"when serializing list item 3" or "when serializing MyClass state".
Both the Python and C implementations are updated, along with tests
asserting the expected notes on all error paths.

Author: pablogsal

Lib/pickle.py

@@ -397,6 +397,14 @@ def decode_long(data):
     return int.from_bytes(data, byteorder='little', signed=True)
 
 
+def _T(obj):
+    cls = type(obj)
+    module = cls.__module__
+    if module in (None, 'builtins', '__main__'):
+        return cls.__qualname__
+    return f'{module}.{cls.__qualname__}'
+
+
 _NoValue = object()
 
 # Pickling machinery
@@ -587,18 +595,22 @@ def save(self, obj, save_persistent_id=True):
             self.save_global(obj, rv)
             return
 
-        # Assert that reduce() returned a tuple
-        if not isinstance(rv, tuple):
-            raise PicklingError("%s must return string or tuple" % reduce)
-
-        # Assert that it returned an appropriately sized tuple
-        l = len(rv)
-        if not (2 <= l <= 6):
-            raise PicklingError("Tuple returned by %s must have "
-                                "two to six elements" % reduce)
-
-        # Save the reduce() output and finally memoize the object
-        self.save_reduce(obj=obj, *rv)
+        try:
+            # Assert that reduce() returned a tuple
+            if not isinstance(rv, tuple):
+                raise PicklingError("%s must return string or tuple" % reduce)
+
+            # Assert that it returned an appropriately sized tuple
+            l = len(rv)
+            if not (2 <= l <= 6):
+                raise PicklingError("Tuple returned by %s must have "
+                                    "two to six elements" % reduce)
+
+            # Save the reduce() output and finally memoize the object
+            self.save_reduce(obj=obj, *rv)
+        except BaseException as exc:
+            exc.add_note(f'when serializing {_T(obj)} object')
+            raise
 
     def persistent_id(self, obj):
         # This exists so a subclass can override it
@@ -638,13 +650,25 @@ def save_reduce(self, func, args, state=None, listitems=None,
                 raise PicklingError("args[0] from {} args has the wrong class"
                                     .format(func_name))
             if self.proto >= 4:
-                save(cls)
-                save(args)
-                save(kwargs)
+                try:
+                    save(cls)
+                except BaseException as exc:
+                    exc.add_note(f'when serializing {_T(obj)} class')
+                    raise
+                try:
+                    save(args)
+                    save(kwargs)
+                except BaseException as exc:
+                    exc.add_note(f'when serializing {_T(obj)} __new__ arguments')
+                    raise
                 write(NEWOBJ_EX)
             else:
                 func = partial(cls.__new__, cls, *args, **kwargs)
-                save(func)
+                try:
+                    save(func)
+                except BaseException as exc:
+                    exc.add_note(f'when serializing {_T(obj)} reconstructor')
+                    raise
                 save(())
                 write(REDUCE)
         elif self.proto >= 2 and func_name == "__newobj__":
@@ -682,12 +706,28 @@ def save_reduce(self, func, args, state=None, listitems=None,
                 raise PicklingError(
                     "args[0] from __newobj__ args has the wrong class")
             args = args[1:]
-            save(cls)
-            save(args)
+            try:
+                save(cls)
+            except BaseException as exc:
+                exc.add_note(f'when serializing {_T(obj)} class')
+                raise
+            try:
+                save(args)
+            except BaseException as exc:
+                exc.add_note(f'when serializing {_T(obj)} __new__ arguments')
+                raise
             write(NEWOBJ)
         else:
-            save(func)
-            save(args)
+            try:
+                save(func)
+            except BaseException as exc:
+                exc.add_note(f'when serializing {_T(obj)} reconstructor')
+                raise
+            try:
+                save(args)
+            except BaseException as exc:
+                exc.add_note(f'when serializing {_T(obj)} reconstructor arguments')
+                raise
             write(REDUCE)
 
         if obj is not None:
@@ -705,23 +745,35 @@ def save_reduce(self, func, args, state=None, listitems=None,
         # items and dict items (as (key, value) tuples), or None.
 
         if listitems is not None:
-            self._batch_appends(listitems)
+            self._batch_appends(listitems, obj)
 
         if dictitems is not None:
-            self._batch_setitems(dictitems)
+            self._batch_setitems(dictitems, obj)
 
         if state is not None:
             if state_setter is None:
-                save(state)
+                try:
+                    save(state)
+                except BaseException as exc:
+                    exc.add_note(f'when serializing {_T(obj)} state')
+                    raise
                 write(BUILD)
             else:
                 # If a state_setter is specified, call it instead of load_build
                 # to update obj's with its previous state.
                 # First, push state_setter and its tuple of expected arguments
                 # (obj, state) onto the stack.
-                save(state_setter)
+                try:
+                    save(state_setter)
+                except BaseException as exc:
+                    exc.add_note(f'when serializing {_T(obj)} state setter')
+                    raise
                 save(obj)  # simple BINGET opcode as obj is already memoized.
-                save(state)
+                try:
+                    save(state)
+                except BaseException as exc:
+                    exc.add_note(f'when serializing {_T(obj)} state')
+                    raise
                 write(TUPLE2)
                 # Trigger a state_setter(obj, state) function call.
                 write(REDUCE)
@@ -901,8 +953,13 @@ def save_tuple(self, obj):
         save = self.save
         memo = self.memo
         if n <= 3 and self.proto >= 2:
-            for element in obj:
-                save(element)
+            for i, element in enumerate(obj):
+                try:
+                    save(element)
+                except BaseException as exc:
+                    exc.add_note(
+                        f'when serializing {_T(obj)} item {i}')
+                    raise
             # Subtle.  Same as in the big comment below.
             if id(obj) in memo:
                 get = self.get(memo[id(obj)][0])
@@ -916,8 +973,13 @@ def save_tuple(self, obj):
         # has more than 3 elements.
         write = self.write
         write(MARK)
-        for element in obj:
-            save(element)
+        for i, element in enumerate(obj):
+            try:
+                save(element)
+            except BaseException as exc:
+                exc.add_note(
+                    f'when serializing {_T(obj)} item {i}')
+                raise
 
         if id(obj) in memo:
             # Subtle.  d was not in memo when we entered save_tuple(), so
@@ -947,34 +1009,55 @@ def save_list(self, obj):
             self.write(MARK + LIST)
 
         self.memoize(obj)
-        self._batch_appends(obj)
+        self._batch_appends(obj, obj)
 
     dispatch[list] = save_list
 
     _BATCHSIZE = 1000
 
-    def _batch_appends(self, items):
+    def _batch_appends(self, items, obj=None):
         # Helper to batch up APPENDS sequences
         save = self.save
         write = self.write
 
         if not self.bin:
-            for x in items:
-                save(x)
+            for i, x in enumerate(items):
+                try:
+                    save(x)
+                except BaseException as exc:
+                    if obj is not None:
+                        exc.add_note(
+                            f'when serializing {_T(obj)} item {i}')
+                    raise
                 write(APPEND)
             return
 
+        i = 0
         it = iter(items)
         while True:
             tmp = list(islice(it, self._BATCHSIZE))
             n = len(tmp)
             if n > 1:
                 write(MARK)
                 for x in tmp:
-                    save(x)
+                    try:
+                        save(x)
+                    except BaseException as exc:
+                        if obj is not None:
+                            exc.add_note(
+                                f'when serializing {_T(obj)} item {i}')
+                        raise
+                    i += 1
                 write(APPENDS)
             elif n:
-                save(tmp[0])
+                try:
+                    save(tmp[0])
+                except BaseException as exc:
+                    if obj is not None:
+                        exc.add_note(
+                            f'when serializing {_T(obj)} item {i}')
+                    raise
+                i += 1
                 write(APPEND)
             # else tmp is empty, and we're done
             if n < self._BATCHSIZE:
@@ -987,19 +1070,25 @@ def save_dict(self, obj):
             self.write(MARK + DICT)
 
         self.memoize(obj)
-        self._batch_setitems(obj.items())
+        self._batch_setitems(obj.items(), obj)
 
     dispatch[dict] = save_dict
 
-    def _batch_setitems(self, items):
+    def _batch_setitems(self, items, obj=None):
         # Helper to batch up SETITEMS sequences; proto >= 1 only
         save = self.save
         write = self.write
 
         if not self.bin:
             for k, v in items:
                 save(k)
-                save(v)
+                try:
+                    save(v)
+                except BaseException as exc:
+                    if obj is not None:
+                        exc.add_note(
+                            f'when serializing {_T(obj)} item {k!r}')
+                    raise
                 write(SETITEM)
             return
 
@@ -1011,12 +1100,24 @@ def _batch_setitems(self, items):
                 write(MARK)
                 for k, v in tmp:
                     save(k)
-                    save(v)
+                    try:
+                        save(v)
+                    except BaseException as exc:
+                        if obj is not None:
+                            exc.add_note(
+                                f'when serializing {_T(obj)} item {k!r}')
+                        raise
                 write(SETITEMS)
             elif n:
                 k, v = tmp[0]
                 save(k)
-                save(v)
+                try:
+                    save(v)
+                except BaseException as exc:
+                    if obj is not None:
+                        exc.add_note(
+                            f'when serializing {_T(obj)} item {k!r}')
+                    raise
                 write(SETITEM)
             # else tmp is empty, and we're done
             if n < self._BATCHSIZE:
@@ -1040,7 +1141,12 @@ def save_set(self, obj):
             if n > 0:
                 write(MARK)
                 for item in batch:
-                    save(item)
+                    try:
+                        save(item)
+                    except BaseException as exc:
+                        exc.add_note(
+                            f'when serializing {_T(obj)} element')
+                        raise
                 write(ADDITEMS)
             if n < self._BATCHSIZE:
                 return
@@ -1056,7 +1162,12 @@ def save_frozenset(self, obj):
 
         write(MARK)
         for item in obj:
-            save(item)
+            try:
+                save(item)
+            except BaseException as exc:
+                exc.add_note(
+                    f'when serializing {_T(obj)} element')
+                raise
 
         if id(obj) in self.memo:
             # If the object is already in the memo, this means it is