Align mapping collection and call unpacking with Python 3.15 (#9302)

* Align mapping collection and call unpacking with Python 3.15

Preserve iterator exceptions and exact-list results when collecting mapping
methods, defer callable formatting, and retain exact-dict stored hashes.
Synchronize complete test_extcall with CPython v3.15.0.

Assisted-by: Codex:unknown

* Update pickle and pickletools libraries for Python 3.15

Synchronize the six coherent serialization modules with CPython v3.15.0
through update_lib, retaining existing RustPython adaptations.

Assisted-by: Codex:unknown

Lib/pickle.py

@@ -17,7 +17,6 @@ Functions:
 
 Misc variables:
 
-    __version__
     format_version
     compatible_formats
 
@@ -190,6 +189,11 @@ READONLY_BUFFER  = b'\x98'  # make top of stack readonly
 __all__.extend(x for x in dir() if x.isupper() and not x.startswith('_'))
 
 
+# Data larger than this will be read in chunks, to prevent extreme
+# overallocation.
+_MIN_READ_BUF_SIZE = (1 << 20)
+
+
 class _Framer:
 
     _FRAME_SIZE_MIN = 4
@@ -288,7 +292,7 @@ class _Unframer:
                     "pickle exhausted before end of frame")
             return data
         else:
-            return self.file_read(n)
+            return self._chunked_file_read(n)
 
     def readline(self):
         if self.current_frame:
@@ -303,11 +307,23 @@ class _Unframer:
         else:
             return self.file_readline()
 
+    def _chunked_file_read(self, size):
+        cursize = min(size, _MIN_READ_BUF_SIZE)
+        b = self.file_read(cursize)
+        while cursize < size and len(b) == cursize:
+            delta = min(cursize, size - cursize)
+            b += self.file_read(delta)
+            cursize += delta
+        return b
+
     def load_frame(self, frame_size):
         if self.current_frame and self.current_frame.read() != b'':
             raise UnpicklingError(
                 "beginning of a new frame before end of current frame")
-        self.current_frame = io.BytesIO(self.file_read(frame_size))
+        data = self._chunked_file_read(frame_size)
+        if len(data) < frame_size:
+            raise EOFError
+        self.current_frame = io.BytesIO(data)
 
 
 # Tools used for pickling.
@@ -1147,13 +1163,23 @@ class _Pickler:
 
     def save_global(self, obj, name=None):
         write = self.write
-        memo = self.memo
 
         if name is None:
             name = getattr(obj, '__qualname__', None)
             if name is None:
                 name = obj.__name__
 
+        if '.__' in name:
+            # Mangle names of private attributes.
+            dotted_path = name.split('.')
+            for i, subpath in enumerate(dotted_path):
+                if i and subpath.startswith('__') and not subpath.endswith('__'):
+                    prev = prev.lstrip('_')
+                    if prev:
+                        dotted_path[i] = f"_{prev.lstrip('_')}{subpath}"
+                prev = subpath
+            name = '.'.join(dotted_path)
+
         module_name = whichmodule(obj, name)
         if self.proto >= 2:
             code = _extension_registry.get((module_name, name), _NoValue)
@@ -1491,12 +1517,17 @@ class _Unpickler:
     dispatch[BINBYTES8[0]] = load_binbytes8
 
     def load_bytearray8(self):
-        len, = unpack('<Q', self.read(8))
-        if len > maxsize:
+        size, = unpack('<Q', self.read(8))
+        if size > maxsize:
             raise UnpicklingError("BYTEARRAY8 exceeds system's maximum size "
                                   "of %d bytes" % maxsize)
-        b = bytearray(len)
-        self.readinto(b)
+        cursize = min(size, _MIN_READ_BUF_SIZE)
+        b = bytearray(cursize)
+        if self.readinto(b) == cursize:
+            while cursize < size and len(b) == cursize:
+                delta = min(cursize, size - cursize)
+                b += self.read(delta)
+                cursize += delta
         self.append(b)
     dispatch[BYTEARRAY8[0]] = load_bytearray8
 
@@ -1729,7 +1760,7 @@ class _Unpickler:
         i = self.read(1)[0]
         try:
             self.append(self.memo[i])
-        except KeyError as exc:
+        except KeyError:
             msg = f'Memo value not found at index {i}'
             raise UnpicklingError(msg) from None
     dispatch[BINGET[0]] = load_binget
@@ -1738,7 +1769,7 @@ class _Unpickler:
         i, = unpack('<I', self.read(4))
         try:
             self.append(self.memo[i])
-        except KeyError as exc:
+        except KeyError:
             msg = f'Memo value not found at index {i}'
             raise UnpicklingError(msg) from None
     dispatch[LONG_BINGET[0]] = load_long_binget
@@ -1906,7 +1937,6 @@ def _main(args=None):
     import pprint
     parser = argparse.ArgumentParser(
         description='display contents of the pickle files',
-        color=True,
     )
     parser.add_argument(
         'pickle_file',

Lib/pickletools.py

@@ -16,6 +16,8 @@ import pickle
 import re
 import sys
 
+lazy from _colorize import decolor, get_theme
+
 __all__ = ['dis', 'genops', 'optimize']
 
 bytes_types = pickle.bytes_types
@@ -2209,6 +2211,32 @@ for i, d in enumerate(opcodes):
     name2i[d.name] = i
     code2i[d.code] = i
 
+# Group opcode names into categories for colourised CLI output.
+_opcode_categories = frozendict(
+    op_call=frozenset({
+        "BUILD", "EXT1", "EXT2", "EXT4", "GLOBAL", "INST", "NEWOBJ",
+        "NEWOBJ_EX", "OBJ", "REDUCE", "STACK_GLOBAL",
+    }),
+    op_container=frozenset({
+        "ADDITEMS", "APPEND", "APPENDS", "DICT", "EMPTY_DICT", "EMPTY_LIST",
+        "EMPTY_SET", "EMPTY_TUPLE", "FROZENSET", "LIST", "SETITEM",
+        "SETITEMS", "TUPLE", "TUPLE1", "TUPLE2", "TUPLE3",
+    }),
+    op_memo=frozenset({
+        "BINGET", "BINPUT", "GET", "LONG_BINGET", "LONG_BINPUT", "MEMOIZE",
+        "PUT",
+    }),
+    op_meta=frozenset({"BINPERSID", "FRAME", "MARK", "PERSID", "PROTO"}),
+    op_stack=frozenset({"DUP", "POP", "POP_MARK", "STOP"}),
+)
+_opcode_color_attr = frozendict({
+    name: attr
+    for attr, names in _opcode_categories.items()
+    for name in names
+})
+assert _opcode_color_attr.keys() <= name2i.keys(), (
+    f"unknown opcodes: {_opcode_color_attr.keys() - name2i.keys()}"
+)
 del name2i, code2i, i, d
 
 ##############################################################################
@@ -2443,13 +2471,19 @@ def dis(pickle, out=None, memo=None, indentlevel=4, annotate=0):
     indentchunk = ' ' * indentlevel
     errormsg = None
     annocol = annotate  # column hint for annotations
+    t = get_theme(tty_file=out).pickletools
     for opcode, arg, pos in genops(pickle):
         if pos is not None:
-            print("%5d:" % pos, end=' ', file=out)
+            print(f"{t.position}{pos:5d}:{t.reset}", end=' ', file=out)
 
-        line = "%-4s %s%s" % (repr(opcode.code)[1:-1],
-                              indentchunk * len(markstack),
-                              opcode.name)
+        attr = _opcode_color_attr.get(opcode.name)
+        opcode_color = getattr(t, attr) if attr else ""
+        opcode_reset = t.reset if attr else ""
+        line = (
+            f"{t.opcode_code}{repr(opcode.code)[1:-1]:<4}{t.reset} "
+            f"{indentchunk * len(markstack)}"
+            f"{opcode_color}{opcode.name}{opcode_reset}"
+        )
 
         maxproto = max(maxproto, opcode.proto)
         before = opcode.stack_before    # don't mutate
@@ -2510,18 +2544,26 @@ def dis(pickle, out=None, memo=None, indentlevel=4, annotate=0):
             line += ' ' * (10 - len(opcode.name))
             if arg is not None:
                 if opcode.name in ("STRING", "BINSTRING", "SHORT_BINSTRING"):
-                    line += ' ' + ascii(arg)
+                    arg_text = ascii(arg)
                 else:
-                    line += ' ' + repr(arg)
+                    arg_text = repr(arg)
+                arg_color = (
+                    t.arg_number
+                    if isinstance(arg, (int, float))
+                    else t.arg_string
+                )
+                line += f" {arg_color}{arg_text}{t.reset}"
             if markmsg:
-                line += ' ' + markmsg
+                line += f" {t.mark}{markmsg}{t.reset}"
         if annotate:
-            line += ' ' * (annocol - len(line))
+            visible_len = len(decolor(line))
+            line += ' ' * (annocol - visible_len)
             # make a mild effort to align annotations
-            annocol = len(line)
+            annocol = max(visible_len, annocol)
             if annocol > 50:
                 annocol = annotate
-            line += ' ' + opcode.doc.split('\n', 1)[0]
+            doc = opcode.doc.split('\n', 1)[0]
+            line += f" {t.annotation}{doc}{t.reset}"
         print(line, file=out)
 
         if errormsg:
@@ -2541,7 +2583,11 @@ def dis(pickle, out=None, memo=None, indentlevel=4, annotate=0):
 
         stack.extend(after)
 
-    print("highest protocol among opcodes =", maxproto, file=out)
+    print(
+        "highest protocol among opcodes =",
+        f"{t.proto}{maxproto}{t.reset}",
+        file=out,
+    )
     if stack:
         raise ValueError("stack not empty after STOP: %r" % stack)
 
@@ -2839,12 +2885,9 @@ __test__ = {'disassembler_test': _dis_test,
            }
 
 
-if __name__ == "__main__":
+def _main(args=None):
     import argparse
-    parser = argparse.ArgumentParser(
-        description='disassemble one or more pickle files',
-        color=True,
-    )
+    parser = argparse.ArgumentParser(description='disassemble one or more pickle files')
     parser.add_argument(
         'pickle_file',
         nargs='+', help='the pickle file')
@@ -2856,7 +2899,7 @@ if __name__ == "__main__":
         help='preserve memo between disassemblies')
     parser.add_argument(
         '-l', '--indentlevel', default=4, type=int,
-        help='the number of blanks by which to indent a new MARK level')
+        help='the number of blanks by which to indent a new `MARK` level')
     parser.add_argument(
         '-a', '--annotate',  action='store_true',
         help='annotate each line with a short opcode description')
@@ -2864,7 +2907,7 @@ if __name__ == "__main__":
         '-p', '--preamble', default="==> {name} <==",
         help='if more than one pickle file is specified, print this before'
         ' each disassembly')
-    args = parser.parse_args()
+    args = parser.parse_args(args)
     annotate = 30 if args.annotate else 0
     memo = {} if args.memo else None
     if args.output is None:
@@ -2885,3 +2928,7 @@ if __name__ == "__main__":
     finally:
         if output is not sys.stdout:
             output.close()
+
+
+if __name__ == "__main__":
+    _main()

Lib/test/picklecommon.py

@@ -263,6 +263,17 @@ myclasses = [MyInt, MyLong, MyFloat,
              MyStr, MyUnicode,
              MyTuple, MyList, MyDict, MySet, MyFrozenSet]
 
+try:
+    frozendict
+except NameError:
+    # Python 3.14 and older
+    pass
+else:
+    class MyFrozenDict(dict):
+        sample = frozendict({"a": 1, "b": 2})
+    myclasses.append(MyFrozenDict)
+
+
 # For test_newobj_overridden_new
 class MyIntWithNew(int):
     def __new__(cls, value):
@@ -388,3 +399,48 @@ class PyMethodsTest:
 class Subclass(tuple):
     class Nested(str):
         pass
+
+# For test_private_methods
+class PrivateMethods:
+    def __init__(self, value):
+        self.value = value
+
+    def __private_method(self):
+        return self.value
+
+    def get_method(self):
+        return self.__private_method
+
+    @classmethod
+    def get_unbound_method(cls):
+        return cls.__private_method
+
+    @classmethod
+    def __private_classmethod(cls):
+        return 43
+
+    @classmethod
+    def get_classmethod(cls):
+        return cls.__private_classmethod
+
+    @staticmethod
+    def __private_staticmethod():
+        return 44
+
+    @classmethod
+    def get_staticmethod(cls):
+        return cls.__private_staticmethod
+
+# For test_private_nested_classes
+class PrivateNestedClasses:
+    @classmethod
+    def get_nested(cls):
+        return cls.__Nested
+
+    class __Nested:
+        @classmethod
+        def get_nested2(cls):
+            return cls.__Nested2
+
+        class __Nested2:
+            pass

Lib/test/pickletester.py

@@ -79,6 +79,15 @@ def count_opcode(code, pickle):
 def identity(x):
     return x
 
+def itersize(start, stop):
+    # Produce geometrical increasing sequence from start to stop
+    # (inclusively) for tests.
+    size = start
+    while size < stop:
+        yield size
+        size <<= 1
+    yield stop
+
 
 class UnseekableIO(io.BytesIO):
     def peek(self, *args):
@@ -751,9 +760,8 @@ class AbstractUnpickleTests:
                 self.assertEqual(getattr(obj, slot, None),
                                  getattr(objcopy, slot, None), msg=msg)
 
-    def check_unpickling_error(self, errors, data):
-        with self.subTest(data=data), \
-             self.assertRaises(errors):
+    def check_unpickling_error_strict(self, errors, data):
+        with self.assertRaises(errors):
             try:
                 self.loads(data)
             except BaseException as exc:
@@ -762,6 +770,10 @@ class AbstractUnpickleTests:
                           (data, exc.__class__.__name__, exc))
                 raise
 
+    def check_unpickling_error(self, errors, data):
+        with self.subTest(data=data):
+            self.check_unpickling_error_strict(errors, data)
+
     def test_load_from_data0(self):
         self.assert_is_copy(self._testdata, self.loads(DATA0))
 
@@ -910,6 +922,16 @@ class AbstractUnpickleTests:
         self.assertIs(self.loads(b'I01\n.'), True)
         self.assertIs(self.loads(b'I00\n.'), False)
 
+    def test_issue135241(self):
+        # C implementation should check for hardcoded values 00 and 01
+        # when getting booleans from the INT opcode. Doing a str comparison
+        # to bypass truthy/falsy comparisons. These payloads should return
+        # 0, not False.
+        out1 = self.loads(b'I+0\n.')
+        self.assertEqual(str(out1), '0')
+        out2 = self.loads(b'I 0\n.')
+        self.assertEqual(str(out2), '0')
+
     def test_zero_padded_integers(self):
         self.assertEqual(self.loads(b'I010\n.'), 10)
         self.assertEqual(self.loads(b'I-010\n.'), -10)
@@ -1038,6 +1060,155 @@ class AbstractUnpickleTests:
         dumped = b'\x80\x03X\x01\x00\x00\x00ar\xff\xff\xff\xff.'
         self.check_unpickling_error(ValueError, dumped)
 
+    def test_too_large_put(self):
+        # Test that PUT with large id does not cause allocation of
+        # too large memo table. The C implementation uses a dict-based memo
+        # for sparse indices (when idx > memo_len * 2) instead of allocating
+        # a massive array. This test verifies large sparse indices work without
+        # causing memory exhaustion.
+        #
+        # The following simple pickle creates an empty list, memoizes it
+        # using a large index, then loads it back on the stack, builds
+        # a tuple containing 2 identical empty lists and returns it.
+        data = lambda n: (b'((lp' + str(n).encode() + b'\n' +
+                          b'g' + str(n).encode() + b'\nt.')
+        #    0: (    MARK
+        #    1: (        MARK
+        #    2: l            LIST       (MARK at 1)
+        #    3: p        PUT        1000000000000
+        #   18: g        GET        1000000000000
+        #   33: t        TUPLE      (MARK at 0)
+        #   34: .    STOP
+        for idx in [10**6, 10**9, 10**12]:
+            if idx > sys.maxsize:
+                continue
+            self.assertEqual(self.loads(data(idx)), ([],)*2)
+
+    def test_too_large_long_binput(self):
+        # Test that LONG_BINPUT with large id does not cause allocation of
+        # too large memo table. The C implementation uses a dict-based memo
+        # for sparse indices (when idx > memo_len * 2) instead of allocating
+        # a massive array. This test verifies large sparse indices work without
+        # causing memory exhaustion.
+        #
+        # The following simple pickle creates an empty list, memoizes it
+        # using a large index, then loads it back on the stack, builds
+        # a tuple containing 2 identical empty lists and returns it.
+        data = lambda n: (b'(]r' + struct.pack('<I', n) +
+                          b'j' + struct.pack('<I', n) + b't.')
+        #    0: (    MARK
+        #    1: ]        EMPTY_LIST
+        #    2: r        LONG_BINPUT 4294967295
+        #    7: j        LONG_BINGET 4294967295
+        #   12: t        TUPLE      (MARK at 0)
+        #   13: .    STOP
+        for idx in itersize(1 << 20, min(sys.maxsize, (1 << 32) - 1)):
+            self.assertEqual(self.loads(data(idx)), ([],)*2)
+
+    def _test_truncated_data(self, dumped, expected_error=None):
+        # Test that instructions to read large data without providing
+        # such amount of data do not cause large memory usage.
+        if expected_error is None:
+            expected_error = self.truncated_data_error
+        # BytesIO
+        with self.assertRaisesRegex(*expected_error):
+            self.loads(dumped)
+        if hasattr(self, 'unpickler'):
+            try:
+                with open(TESTFN, 'wb') as f:
+                    f.write(dumped)
+                # buffered file
+                with open(TESTFN, 'rb') as f:
+                    u = self.unpickler(f)
+                    with self.assertRaisesRegex(*expected_error):
+                        u.load()
+                # unbuffered file
+                with open(TESTFN, 'rb', buffering=0) as f:
+                    u = self.unpickler(f)
+                    with self.assertRaisesRegex(*expected_error):
+                        u.load()
+            finally:
+                os_helper.unlink(TESTFN)
+
+    def test_truncated_large_binstring(self):
+        data = lambda size: b'T' + struct.pack('<I', size) + b'.' * 5
+        #    0: T    BINSTRING  '....'
+        #    9: .    STOP
+        self.assertEqual(self.loads(data(4)), '....') # self-testing
+        for size in itersize(1 << 10, min(sys.maxsize - 5, (1 << 31) - 1)):
+            self._test_truncated_data(data(size))
+        self._test_truncated_data(data(1 << 31),
+            (pickle.UnpicklingError, 'truncated|exceeds|negative byte count'))
+
+    def test_truncated_large_binunicode(self):
+        data = lambda size: b'X' + struct.pack('<I', size) + b'.' * 5
+        #    0: X    BINUNICODE '....'
+        #    9: .    STOP
+        self.assertEqual(self.loads(data(4)), '....') # self-testing
+        for size in itersize(1 << 10, min(sys.maxsize - 5, (1 << 32) - 1)):
+            self._test_truncated_data(data(size))
+
+    def test_truncated_large_binbytes(self):
+        data = lambda size: b'B' + struct.pack('<I', size) + b'.' * 5
+        #    0: B    BINBYTES   b'....'
+        #    9: .    STOP
+        self.assertEqual(self.loads(data(4)), b'....') # self-testing
+        for size in itersize(1 << 10, min(sys.maxsize, 1 << 31)):
+            self._test_truncated_data(data(size))
+
+    def test_truncated_large_long4(self):
+        data = lambda size: b'\x8b' + struct.pack('<I', size) + b'.' * 5
+        #    0: \x8b LONG4      0x2e2e2e2e
+        #    9: .    STOP
+        self.assertEqual(self.loads(data(4)), 0x2e2e2e2e) # self-testing
+        for size in itersize(1 << 10, min(sys.maxsize - 5, (1 << 31) - 1)):
+            self._test_truncated_data(data(size))
+        self._test_truncated_data(data(1 << 31),
+            (pickle.UnpicklingError, 'LONG pickle has negative byte count'))
+
+    def test_truncated_large_frame(self):
+        data = lambda size: b'\x95' + struct.pack('<Q', size) + b'N.'
+        #    0: \x95 FRAME      2
+        #    9: N    NONE
+        #   10: .    STOP
+        self.assertIsNone(self.loads(data(2))) # self-testing
+        for size in itersize(1 << 10, sys.maxsize - 9):
+            self._test_truncated_data(data(size))
+        if sys.maxsize + 1 < 1 << 64:
+            self._test_truncated_data(data(sys.maxsize + 1),
+                ((OverflowError, ValueError),
+                 'FRAME length exceeds|frame size > sys.maxsize'))
+
+    def test_truncated_large_binunicode8(self):
+        data = lambda size: b'\x8d' + struct.pack('<Q', size) + b'.' * 5
+        #    0: \x8d BINUNICODE8 '....'
+        #   13: .    STOP
+        self.assertEqual(self.loads(data(4)), '....') # self-testing
+        for size in itersize(1 << 10, sys.maxsize - 9):
+            self._test_truncated_data(data(size))
+        if sys.maxsize + 1 < 1 << 64:
+            self._test_truncated_data(data(sys.maxsize + 1), self.size_overflow_error)
+
+    def test_truncated_large_binbytes8(self):
+        data = lambda size: b'\x8e' + struct.pack('<Q', size) + b'.' * 5
+        #    0: \x8e BINBYTES8  b'....'
+        #   13: .    STOP
+        self.assertEqual(self.loads(data(4)), b'....') # self-testing
+        for size in itersize(1 << 10, sys.maxsize):
+            self._test_truncated_data(data(size))
+        if sys.maxsize + 1 < 1 << 64:
+            self._test_truncated_data(data(sys.maxsize + 1), self.size_overflow_error)
+
+    def test_truncated_large_bytearray8(self):
+        data = lambda size: b'\x96' + struct.pack('<Q', size) + b'.' * 5
+        #    0: \x96 BYTEARRAY8 bytearray(b'....')
+        #   13: .    STOP
+        self.assertEqual(self.loads(data(4)), bytearray(b'....')) # self-testing
+        for size in itersize(1 << 10, sys.maxsize):
+            self._test_truncated_data(data(size))
+        if sys.maxsize + 1 < 1 << 64:
+            self._test_truncated_data(data(sys.maxsize + 1), self.size_overflow_error)
+
     def test_badly_escaped_string(self):
         self.check_unpickling_error(ValueError, b"S'\\'\n.")
 
@@ -1357,6 +1528,29 @@ class AbstractUnpickleTests:
         # bad hashable dict key
         self.check_unpickling_error(CustomError, base + b'}c__main__\nBadKey1\n)\x81Nsb.')
 
+    def test_bad_types(self):
+        # APPEND
+        self.assertEqual(self.loads(b']Na.'), [None])
+        self.check_unpickling_error(AttributeError, b'NNa.')  # non-list
+        # APPENDS
+        self.assertEqual(self.loads(b'](Ne.'), [None])
+        self.check_unpickling_error(AttributeError, b'N(Ne.')  # non-list
+        self.check_unpickling_error(AttributeError, b'N(e.')
+        # SETITEM
+        self.assertEqual(self.loads(b'}NNs.'), {None: None})
+        self.check_unpickling_error(TypeError, b'NNNs.')  # non-dict
+        self.check_unpickling_error(TypeError, b'}]Ns.')  # non-hashable key
+        # SETITEMS
+        self.assertEqual(self.loads(b'}(NNu.'), {None: None})
+        self.check_unpickling_error(TypeError, b'N(NNu.')  # non-dict
+        self.assertEqual(self.loads(b'N(u.'), None)  # no validation for empty items
+        self.check_unpickling_error(TypeError, b'}(]Nu.')  # non-hashable key
+        # ADDITEMS
+        self.assertEqual(self.loads(b'\x8f(N\x90.'), {None})
+        self.check_unpickling_error(AttributeError, b'N(N\x90.')  # non-set
+        self.check_unpickling_error(AttributeError, b'N(\x90.')
+        self.check_unpickling_error(TypeError, b'\x8f(]\x90.')  # non-hashable element
+
     def test_bad_stack(self):
         badpickles = [
             b'.',                       # STOP
@@ -1812,7 +2006,7 @@ class AbstractPicklingErrorTests:
                     with self.assertRaises(TypeError) as cm:
                         self.dumps(obj, proto)
                     self.assertEqual(str(cm.exception),
-                        'functools.partial() argument after ** must be a mapping, not list')
+                        'Value after ** must be a mapping, not list')
                     self.assertEqual(cm.exception.__notes__, [
                         f'when serializing {REX.__module__}.REX object'])
         else:
@@ -2689,11 +2883,13 @@ class AbstractPickleTests:
             self.assertEqual(list(x[0].attr.keys()), [1])
             self.assertIs(x[0].attr[1], x)
 
-    def _test_recursive_collection_and_inst(self, factory, oldminproto=None):
+    def _test_recursive_collection_and_inst(self, factory, oldminproto=None,
+                                            minprotocol=0):
         if self.py_version < (3, 0):
             self.skipTest('"classic" classes are not interoperable with Python 2')
         # Mutable object containing a collection containing the original
         # object.
+        protocols = range(minprotocol, pickle.HIGHEST_PROTOCOL + 1)
         o = Object()
         o.attr = factory([o])
         t = type(o.attr)
@@ -2733,6 +2929,11 @@ class AbstractPickleTests:
     def test_recursive_dict_and_inst(self):
         self._test_recursive_collection_and_inst(dict.fromkeys, oldminproto=0)
 
+    def test_recursive_frozendict_and_inst(self):
+        if self.py_version < (3, 15):
+            self.skipTest('need frozendict')
+        self._test_recursive_collection_and_inst(frozendict.fromkeys, minprotocol=2)
+
     def test_recursive_set_and_inst(self):
         self._test_recursive_collection_and_inst(set)
 
@@ -2754,6 +2955,50 @@ class AbstractPickleTests:
     def test_recursive_frozenset_subclass_and_inst(self):
         self._test_recursive_collection_and_inst(MyFrozenSet)
 
+    def _test_recursive_collection_in_key(self, factory, minprotocol=0):
+        protocols = range(minprotocol, pickle.HIGHEST_PROTOCOL + 1)
+        key = Object()
+        o = factory({key: 1})
+        key.attr = o
+        for proto in protocols:
+            with self.subTest(proto=proto):
+                s = self.dumps(o, proto)
+                x = self.loads(s)
+                keys = list(x.keys())
+                self.assertEqual(len(keys), 1)
+                self.assertIs(keys[0].attr, x)
+
+    def test_recursive_frozendict_in_key(self):
+        if self.py_version < (3, 15):
+            self.skipTest('need frozendict')
+        self._test_recursive_collection_in_key(frozendict, minprotocol=2)
+
+    def test_recursive_frozendict_subclass_in_key(self):
+        if self.py_version < (3, 15):
+            self.skipTest('need frozendict')
+        self._test_recursive_collection_in_key(MyFrozenDict)
+
+    def _test_recursive_collection_in_value(self, factory, minprotocol=0):
+        protocols = range(minprotocol, pickle.HIGHEST_PROTOCOL + 1)
+        o = factory(key=[])
+        o['key'].append(o)
+        for proto in protocols:
+            with self.subTest(proto=proto):
+                s = self.dumps(o, proto)
+                x = self.loads(s)
+                self.assertEqual(len(x['key']), 1)
+                self.assertIs(x['key'][0], x)
+
+    def test_recursive_frozendict_in_value(self):
+        if self.py_version < (3, 15):
+            self.skipTest('need frozendict')
+        self._test_recursive_collection_in_value(frozendict, minprotocol=2)
+
+    def test_recursive_frozendict_subclass_in_value(self):
+        if self.py_version < (3, 15):
+            self.skipTest('need frozendict')
+        self._test_recursive_collection_in_value(MyFrozenDict)
+
     def test_recursive_inst_state(self):
         # Mutable object containing itself.
         y = REX_state()
@@ -3085,6 +3330,8 @@ class AbstractPickleTests:
             'bytes': (3, 0),
             'BuiltinImporter': (3, 3),
             'str': (3, 4),  # not interoperable with Python < 3.4
+            'frozendict': (3, 15),
+            'sentinel': (3, 15),
         }
         for t in builtins.__dict__.values():
             if isinstance(t, type) and not issubclass(t, BaseException):
@@ -3120,6 +3367,7 @@ class AbstractPickleTests:
             'ExceptionGroup': (3, 11),
             '_IncompleteInputError': (3, 13),
             'PythonFinalizationError': (3, 13),
+            'ImportCycleError': (3, 15),
         }
         for t in builtins.__dict__.values():
             if isinstance(t, type) and issubclass(t, BaseException):
@@ -3146,6 +3394,7 @@ class AbstractPickleTests:
             'breakpoint': (3, 7),
             'aiter': (3, 10),
             'anext': (3, 10),
+            '__lazy_import__': (3, 15),
         }
         for t in builtins.__dict__.values():
             if isinstance(t, types.BuiltinFunctionType):
@@ -3286,6 +3535,8 @@ class AbstractPickleTests:
                         self.skipTest('int and str subclasses are not interoperable with Python 2')
                     if (3, 0) <= self.py_version < (3, 4) and proto < 2 and C in (MyStr, MyUnicode):
                         self.skipTest('str subclasses are not interoperable with Python < 3.4')
+                    if self.py_version < (3, 15) and C == MyFrozenDict:
+                        self.skipTest('frozendict is not available on Python < 3.15')
                     B = C.__base__
                     x = C(C.sample)
                     x.foo = 42
@@ -3307,6 +3558,8 @@ class AbstractPickleTests:
                 with self.subTest(proto=proto, C=C):
                     if self.py_version < (3, 4) and proto < 3 and C in (MyStr, MyUnicode):
                         self.skipTest('str subclasses are not interoperable with Python < 3.4')
+                    if self.py_version < (3, 15) and C == MyFrozenDict:
+                        self.skipTest('frozendict is not available on Python < 3.15')
                     B = C.__base__
                     x = C(C.sample)
                     x.foo = 42
@@ -4036,6 +4289,33 @@ class AbstractPickleTests:
                 with self.subTest(proto=proto, descr=descr):
                     self.assertRaises(TypeError, self.dumps, descr, proto)
 
+    def test_private_methods(self):
+        if self.py_version < (3, 15):
+            self.skipTest('not supported in Python < 3.15')
+        obj = PrivateMethods(42)
+        for proto in protocols:
+            with self.subTest(proto=proto):
+                unpickled = self.loads(self.dumps(obj.get_method(), proto))
+                self.assertEqual(unpickled(), 42)
+                unpickled = self.loads(self.dumps(obj.get_unbound_method(), proto))
+                self.assertEqual(unpickled(obj), 42)
+                unpickled = self.loads(self.dumps(obj.get_classmethod(), proto))
+                self.assertEqual(unpickled(), 43)
+                unpickled = self.loads(self.dumps(obj.get_staticmethod(), proto))
+                self.assertEqual(unpickled(), 44)
+
+    def test_private_nested_classes(self):
+        if self.py_version < (3, 15):
+            self.skipTest('not supported in Python < 3.15')
+        cls1 = PrivateNestedClasses.get_nested()
+        cls2 = cls1.get_nested2()
+        for proto in protocols:
+            with self.subTest(proto=proto):
+                unpickled = self.loads(self.dumps(cls1, proto))
+                self.assertIs(unpickled, cls1)
+                unpickled = self.loads(self.dumps(cls2, proto))
+                self.assertIs(unpickled, cls2)
+
     def test_object_with_attrs(self):
         obj = Object()
         obj.a = 1

Lib/test/test_extcall.py

@@ -137,7 +137,7 @@ Verify clearing of SF bug #733667
     >>> g(*Nothing())
     Traceback (most recent call last):
       ...
-    TypeError: test.test_extcall.g() argument after * must be an iterable, not Nothing
+    TypeError: Value after * must be an iterable, not Nothing
 
     >>> class Nothing:
     ...     def __len__(self): return 5
@@ -146,7 +146,7 @@ Verify clearing of SF bug #733667
     >>> g(*Nothing())
     Traceback (most recent call last):
       ...
-    TypeError: test.test_extcall.g() argument after * must be an iterable, not Nothing
+    TypeError: Value after * must be an iterable, not Nothing
 
     >>> class Nothing():
     ...     def __len__(self): return 5
@@ -266,7 +266,7 @@ What about willful misconduct?
     >>> h(*h)
     Traceback (most recent call last):
       ...
-    TypeError: test.test_extcall.h() argument after * must be an iterable, not function
+    TypeError: Value after * must be an iterable, not function
 
     >>> h(1, *h)
     Traceback (most recent call last):
@@ -281,55 +281,69 @@ What about willful misconduct?
     >>> dir(*h)
     Traceback (most recent call last):
       ...
-    TypeError: dir() argument after * must be an iterable, not function
+    TypeError: Value after * must be an iterable, not function
 
     >>> nothing = None
     >>> nothing(*h)
     Traceback (most recent call last):
       ...
-    TypeError: None argument after * must be an iterable, \
-not function
+    TypeError: Value after * must be an iterable, not function
 
     >>> h(**h)
     Traceback (most recent call last):
       ...
-    TypeError: test.test_extcall.h() argument after ** must be a mapping, not function
+    TypeError: Value after ** must be a mapping, not function
 
     >>> h(**[])
     Traceback (most recent call last):
       ...
-    TypeError: test.test_extcall.h() argument after ** must be a mapping, not list
+    TypeError: Value after ** must be a mapping, not list
 
     >>> h(a=1, **h)
     Traceback (most recent call last):
       ...
-    TypeError: test.test_extcall.h() argument after ** must be a mapping, not function
+    TypeError: Value after ** must be a mapping, not function
 
     >>> h(a=1, **[])
     Traceback (most recent call last):
       ...
-    TypeError: test.test_extcall.h() argument after ** must be a mapping, not list
+    TypeError: Value after ** must be a mapping, not list
 
     >>> h(**{'a': 1}, **h)
     Traceback (most recent call last):
       ...
-    TypeError: test.test_extcall.h() argument after ** must be a mapping, not function
+    TypeError: Value after ** must be a mapping, not function
 
     >>> h(**{'a': 1}, **[])
     Traceback (most recent call last):
       ...
-    TypeError: test.test_extcall.h() argument after ** must be a mapping, not list
+    TypeError: Value after ** must be a mapping, not list
 
     >>> dir(**h)
     Traceback (most recent call last):
       ...
-    TypeError: dir() argument after ** must be a mapping, not function
+    TypeError: Value after ** must be a mapping, not function
 
     >>> nothing(**h)
     Traceback (most recent call last):
       ...
-    TypeError: None argument after ** must be a mapping, \
-not function
+    TypeError: Value after ** must be a mapping, not function
+
+    >>> class OnlyKeys:
+    ...     def keys(self):
+    ...         return ['key']
+    >>> h(**OnlyKeys())
+    Traceback (most recent call last):
+      ...
+    TypeError: 'OnlyKeys' object is not subscriptable
+
+    >>> class BrokenKeys:
+    ...     def keys(self):
+    ...         return 1
+    >>> h(**BrokenKeys())
+    Traceback (most recent call last):
+      ...
+    TypeError: test.test_extcall.BrokenKeys.keys() must return an iterable, not int
 
     >>> dir(b=1, **{'b': 1})
     Traceback (most recent call last):
@@ -542,18 +556,156 @@ Same with keyword only args:
 
 """
 
+def test_errors_in_iter():
+    """
+    >>> class A:
+    ...     def __iter__(self):
+    ...         raise exc
+    ...
+    >>> def f(*args, **kwargs): pass
+    >>> exc = ZeroDivisionError('some error')
+    >>> f(*A())
+    Traceback (most recent call last):
+      ...
+    ZeroDivisionError: some error
+
+    >>> exc = AttributeError('some error')
+    >>> f(*A())
+    Traceback (most recent call last):
+      ...
+    AttributeError: some error
+
+    >>> exc = TypeError('some error')
+    >>> f(*A())
+    Traceback (most recent call last):
+      ...
+    TypeError: some error
+    """
+
+def test_errors_in_next():
+    """
+    >>> class I:
+    ...     def __iter__(self):
+    ...         return self
+    ...     def __next__(self):
+    ...         raise exc
+    ...
+    >>> class A:
+    ...     def __iter__(self):
+    ...         return I()
+    ...
+
+    >>> def f(*args, **kwargs): pass
+    >>> exc = ZeroDivisionError('some error')
+    >>> f(*A())
+    Traceback (most recent call last):
+      ...
+    ZeroDivisionError: some error
+
+    >>> exc = AttributeError('some error')
+    >>> f(*A())
+    Traceback (most recent call last):
+      ...
+    AttributeError: some error
+
+    >>> exc = TypeError('some error')
+    >>> f(*A())
+    Traceback (most recent call last):
+      ...
+    TypeError: some error
+    """
+
+def test_errors_in_keys():
+    """
+    >>> class D:
+    ...     def keys(self):
+    ...         raise exc
+    ...
+    >>> def f(*args, **kwargs): pass
+    >>> exc = ZeroDivisionError('some error')
+    >>> f(**D())
+    Traceback (most recent call last):
+      ...
+    ZeroDivisionError: some error
+
+    >>> exc = AttributeError('some error')
+    >>> f(**D())
+    Traceback (most recent call last):
+      ...
+    AttributeError: some error
+
+    >>> exc = KeyError('some error')
+    >>> f(**D())
+    Traceback (most recent call last):
+      ...
+    KeyError: 'some error'
+    """
+
+def test_errors_in_keys_next():
+    """
+    >>> class I:
+    ...     def __iter__(self):
+    ...         return self
+    ...     def __next__(self):
+    ...         raise exc
+    ...
+    >>> class D:
+    ...     def keys(self):
+    ...         return I()
+    ...
+    >>> def f(*args, **kwargs): pass
+    >>> exc = ZeroDivisionError('some error')
+    >>> f(**D())
+    Traceback (most recent call last):
+      ...
+    ZeroDivisionError: some error
+
+    >>> exc = AttributeError('some error')
+    >>> f(**D())
+    Traceback (most recent call last):
+      ...
+    AttributeError: some error
+
+    >>> exc = KeyError('some error')
+    >>> f(**D())
+    Traceback (most recent call last):
+      ...
+    KeyError: 'some error'
+    """
+
+def test_errors_in_getitem():
+    """
+    >>> class D:
+    ...     def keys(self):
+    ...         return ['key']
+    ...     def __getitem__(self, key):
+    ...         raise exc
+    ...
+    >>> def f(*args, **kwargs): pass
+    >>> exc = ZeroDivisionError('some error')
+    >>> f(**D())
+    Traceback (most recent call last):
+      ...
+    ZeroDivisionError: some error
+
+    >>> exc = AttributeError('some error')
+    >>> f(**D())
+    Traceback (most recent call last):
+      ...
+    AttributeError: some error
+
+    >>> exc = KeyError('some error')
+    >>> f(**D())
+    Traceback (most recent call last):
+      ...
+    KeyError: 'some error'
+    """
+
 import doctest
 import unittest
 
-EXPECTED_FAILURE = doctest.register_optionflag('EXPECTED_FAILURE') # TODO: RUSTPYTHON
-class CustomOutputChecker(doctest.OutputChecker): # TODO: RUSTPYTHON
-    def check_output(self, want, got, optionflags): # TODO: RUSTPYTHON
-        if optionflags & EXPECTED_FAILURE: # TODO: RUSTPYTHON
-            return not super().check_output(want, got, optionflags) # TODO: RUSTPYTHON
-        return super().check_output(want, got, optionflags) # TODO: RUSTPYTHON
-
 def load_tests(loader, tests, pattern):
-    tests.addTest(doctest.DocTestSuite(checker=CustomOutputChecker())) # TODO: RUSTPYTHON
+    tests.addTest(doctest.DocTestSuite())
     return tests
 
 

Lib/test/test_pickle.py

@@ -52,8 +52,6 @@ class PyPickleTests(AbstractPickleModuleTests, unittest.TestCase):
     Unpickler = pickle._Unpickler
 
 
-
-
 class PyUnpicklerTests(AbstractUnpickleTests, unittest.TestCase):
 
     unpickler = pickle._Unpickler
@@ -61,6 +59,8 @@ class PyUnpicklerTests(AbstractUnpickleTests, unittest.TestCase):
     truncated_errors = (pickle.UnpicklingError, EOFError,
                         AttributeError, ValueError,
                         struct.error, IndexError, ImportError)
+    truncated_data_error = (EOFError, '')
+    size_overflow_error = (pickle.UnpicklingError, 'exceeds')
 
     def loads(self, buf, **kwds):
         f = io.BytesIO(buf)
@@ -80,9 +80,6 @@ class PyPicklingErrorTests(AbstractPicklingErrorTests, unittest.TestCase):
         return bytes(f.read())
 
 
-
-
-
 class PyPicklerTests(AbstractPickleTests, unittest.TestCase):
 
     pickler = pickle._Pickler
@@ -100,9 +97,6 @@ class PyPicklerTests(AbstractPickleTests, unittest.TestCase):
         u = self.unpickler(f, **kwds)
         return u.load()
 
-
-
-
     @unittest.expectedFailure  # TODO: RUSTPYTHON
     def test_c_methods(self):
         return super().test_c_methods()
@@ -112,9 +106,6 @@ class PyPicklerTests(AbstractPickleTests, unittest.TestCase):
         return super().test_deep_nested_struct_frozendict()
 
 
-
-
-
 class InMemoryPickleTests(AbstractPickleTests, AbstractUnpickleTests,
                           BigmemPickleTests, unittest.TestCase):
 
@@ -122,6 +113,8 @@ class InMemoryPickleTests(AbstractPickleTests, AbstractUnpickleTests,
     truncated_errors = (pickle.UnpicklingError, EOFError,
                         AttributeError, ValueError,
                         struct.error, IndexError, ImportError)
+    truncated_data_error = ((pickle.UnpicklingError, EOFError), '')
+    size_overflow_error = ((OverflowError, pickle.UnpicklingError), 'exceeds')
 
     def dumps(self, arg, protocol=None, **kwargs):
         return pickle.dumps(arg, protocol, **kwargs)
@@ -133,17 +126,11 @@ class InMemoryPickleTests(AbstractPickleTests, AbstractUnpickleTests,
     test_find_class = None
     test_custom_find_class = None
 
-
-
-
     @unittest.expectedFailure  # TODO: RUSTPYTHON
     def test_c_methods(self):
         return super().test_c_methods()
 
 
-
-
-
 class PersistentPicklerUnpicklerMixin(object):
 
     def dumps(self, arg, proto=None):
@@ -404,6 +391,8 @@ if has_c_implementation:
         unpickler = _pickle.Unpickler
         bad_stack_errors = (pickle.UnpicklingError,)
         truncated_errors = (pickle.UnpicklingError,)
+        truncated_data_error = (pickle.UnpicklingError, 'truncated')
+        size_overflow_error = (OverflowError, 'exceeds')
 
     class CPicklingErrorTests(PyPicklingErrorTests):
         pickler = _pickle.Pickler
@@ -555,7 +544,7 @@ if has_c_implementation:
                 0)  # Write buffer is cleared after every dump().
 
         def test_unpickler(self):
-            basesize = support.calcobjsize('2P2n2P 2P2n2i5P 2P5n8P2n3i')
+            basesize = support.calcobjsize('2P2n3P 2P2n2i5P 2P5n8P2n3i')
             unpickler = _pickle.Unpickler
             P = struct.calcsize('P')  # Size of memo table entry.
             n = struct.calcsize('n')  # Size of mark table entry.

Lib/test/test_pickletools.py

@@ -1,7 +1,11 @@
 import io
+import itertools
 import pickle
 import pickletools
+import tempfile
+import textwrap
 from test import support
+from test.support import os_helper
 from test.pickletester import AbstractPickleTests
 import doctest
 import unittest
@@ -62,17 +66,11 @@ class OptimizedPickleTests(AbstractPickleTests, unittest.TestCase):
         self.assertIs(unpickled2[1], unpickled2[2])
         self.assertNotIn(pickle.BINPUT, pickled2)
 
-
-
-
     @unittest.expectedFailure  # TODO: RUSTPYTHON
     def test_c_methods(self):
         return super().test_c_methods()
 
 
-
-
-
 class SimpleReader:
     def __init__(self, data):
         self.data = data
@@ -166,6 +164,7 @@ class GenopsTests(unittest.TestCase):
             next(it)
 
 
+@support.force_not_colorized_test_class
 class DisTests(unittest.TestCase):
     maxDiff = None
 
@@ -524,6 +523,171 @@ class MiscTestCase(unittest.TestCase):
         support.check__all__(self, pickletools, not_exported=not_exported)
 
 
+@support.force_not_colorized_test_class
+class CommandLineTest(unittest.TestCase):
+    def setUp(self):
+        self.filename = tempfile.mktemp()
+        self.addCleanup(os_helper.unlink, self.filename)
+
+    @staticmethod
+    def text_normalize(string):
+        return textwrap.dedent(string).strip()
+
+    def set_pickle_data(self, data):
+        with open(self.filename, 'wb') as f:
+            pickle.dump(data, f)
+
+    def invoke_pickletools(self, *flags):
+        with (
+            support.captured_stdout() as stdout,
+            support.captured_stderr() as stderr,
+        ):
+            pickletools._main(args=[*flags, self.filename])
+        self.assertEqual(stderr.getvalue(), '')
+        return self.text_normalize(stdout.getvalue())
+
+    def check_output(self, data, expect, *flags):
+        with self.subTest(data=data, flags=flags):
+            self.set_pickle_data(data)
+            res = self.invoke_pickletools(*flags)
+            expect = self.text_normalize(expect)
+            self.assertListEqual(res.splitlines(), expect.splitlines())
+
+    def test_invocation(self):
+        # test various combinations of parameters
+        output_file = tempfile.mktemp()
+        self.addCleanup(os_helper.unlink, output_file)
+        base_flags = [
+            (f'-o={output_file}', f'--output={output_file}'),
+            ('-m', '--memo'),
+            ('-l=2', '--indentlevel=2'),
+            ('-a', '--annotate'),
+            ('-p="Another:"', '--preamble="Another:"'),
+        ]
+        data = { 'a', 'b', 'c' }
+
+        self.set_pickle_data(data)
+
+        for r in range(1, len(base_flags) + 1):
+            for choices in itertools.combinations(base_flags, r=r):
+                for args in itertools.product(*choices):
+                    with self.subTest(args=args[1:]):
+                        self.invoke_pickletools(*args)
+
+    def test_unknown_flag(self):
+        with self.assertRaises(SystemExit):
+            with support.captured_stderr() as stderr:
+                pickletools._main(args=['--unknown'])
+            self.assertStartsWith(stderr.getvalue(), 'usage: ')
+
+    def test_output_flag(self):
+        # test 'python -m pickletools -o/--output'
+        output_file = tempfile.mktemp()
+        self.addCleanup(os_helper.unlink, output_file)
+        data = ('fake_data',)
+        expect = r'''
+            0: \x80 PROTO      5
+            2: \x95 FRAME      15
+           11: \x8c SHORT_BINUNICODE 'fake_data'
+           22: \x94 MEMOIZE    (as 0)
+           23: \x85 TUPLE1
+           24: \x94 MEMOIZE    (as 1)
+           25: .    STOP
+        highest protocol among opcodes = 4
+        '''
+        for flag in [f'-o={output_file}', f'--output={output_file}']:
+            with self.subTest(data=data, flags=flag):
+                self.set_pickle_data(data)
+                res = self.invoke_pickletools(flag)
+                with open(output_file, 'r') as f:
+                    res_from_file = self.text_normalize(f.read())
+                expect = self.text_normalize(expect)
+
+                self.assertListEqual(res.splitlines(), [])
+                self.assertListEqual(res_from_file.splitlines(),
+                                     expect.splitlines())
+
+    def test_memo_flag(self):
+        # test 'python -m pickletools -m/--memo'
+        data = ('fake_data',)
+        expect = r'''
+            0: \x80 PROTO      5
+            2: \x95 FRAME      15
+           11: \x8c SHORT_BINUNICODE 'fake_data'
+           22: \x94 MEMOIZE    (as 0)
+           23: \x85 TUPLE1
+           24: \x94 MEMOIZE    (as 1)
+           25: .    STOP
+        highest protocol among opcodes = 4
+        '''
+        for flag in ['-m', '--memo']:
+            self.check_output(data, expect, flag)
+
+    def test_indentlevel_flag(self):
+        # test 'python -m pickletools -l/--indentlevel'
+        data = ('fake_data',)
+        expect = r'''
+            0: \x80 PROTO      5
+            2: \x95 FRAME      15
+           11: \x8c SHORT_BINUNICODE 'fake_data'
+           22: \x94 MEMOIZE    (as 0)
+           23: \x85 TUPLE1
+           24: \x94 MEMOIZE    (as 1)
+           25: .    STOP
+        highest protocol among opcodes = 4
+        '''
+        for flag in ['-l=2', '--indentlevel=2']:
+            self.check_output(data, expect, flag)
+
+    def test_annotate_flag(self):
+        # test 'python -m pickletools -a/--annotate'
+        data = ('fake_data',)
+        expect = r'''
+            0: \x80 PROTO      5              Protocol version indicator.
+            2: \x95 FRAME      15             Indicate the beginning of a new frame.
+           11: \x8c SHORT_BINUNICODE 'fake_data' Push a Python Unicode string object.
+           22: \x94 MEMOIZE    (as 0)            Store the stack top into the memo.  The stack is not popped.
+           23: \x85 TUPLE1                       Build a one-tuple out of the topmost item on the stack.
+           24: \x94 MEMOIZE    (as 1)            Store the stack top into the memo.  The stack is not popped.
+           25: .    STOP                         Stop the unpickling machine.
+        highest protocol among opcodes = 4
+        '''
+        for flag in ['-a', '--annotate']:
+            self.check_output(data, expect, flag)
+
+    def test_preamble_flag(self):
+        # test 'python -m pickletools -p/--preamble'
+        data = ('fake_data',)
+        expect = r'''
+        Another:
+            0: \x80 PROTO      5
+            2: \x95 FRAME      15
+           11: \x8c SHORT_BINUNICODE 'fake_data'
+           22: \x94 MEMOIZE    (as 0)
+           23: \x85 TUPLE1
+           24: \x94 MEMOIZE    (as 1)
+           25: .    STOP
+        highest protocol among opcodes = 4
+        Another:
+            0: \x80 PROTO      5
+            2: \x95 FRAME      15
+           11: \x8c SHORT_BINUNICODE 'fake_data'
+           22: \x94 MEMOIZE    (as 0)
+           23: \x85 TUPLE1
+           24: \x94 MEMOIZE    (as 1)
+           25: .    STOP
+        highest protocol among opcodes = 4
+        '''
+        for flag in ['-p=Another:', '--preamble=Another:']:
+            with self.subTest(data=data, flags=flag):
+                self.set_pickle_data(data)
+                with support.captured_stdout() as stdout:
+                    pickletools._main(args=[flag, self.filename, self.filename])
+                res = self.text_normalize(stdout.getvalue())
+                expect = self.text_normalize(expect)
+                self.assertListEqual(res.splitlines(), expect.splitlines())
+
+
 def load_tests(loader, tests, pattern):
     from test.support.rustpython import DocTestChecker # TODO: RUSTPYTHON; Remove this
     tests.addTest(doctest.DocTestSuite(pickletools, checker=DocTestChecker())) # XXX: RUSTPYTHON

crates/vm/src/builtins/dict.rs

@@ -424,7 +424,12 @@ impl PyDict {
             }
             Err(e) => return Err(e),
         };
-        let keys = PyIter::try_from_object(vm, keys_method.call((), vm)?)?;
+        let keys = other.mapping_unchecked().collect_method_output(
+            vm.ctx.intern_str("keys"),
+            keys_method.call((), vm)?,
+            vm,
+        )?;
+        let keys = PyIter::try_from_object(vm, keys)?;
         while let PyIterReturn::Return(key) = keys.next(vm)? {
             if !override_existing {
                 let hash = key_hash(&*key, mapping_name, vm)?;

crates/vm/src/frame.rs

@@ -50,7 +50,7 @@ use num_traits::{ToPrimitive, Zero};
 use rustpython_common::atomic::{PyAtomic, Radium};
 use rustpython_common::{
     lock::{OnceCell, PyMutex},
-    wtf8::{Wtf8, Wtf8Buf, wtf8_concat},
+    wtf8::{Wtf8Buf, wtf8_concat},
 };
 use rustpython_compiler_core::{OneIndexed, SourceLocation};
 
@@ -4880,65 +4880,71 @@ impl ExecutingFrame<'_> {
 
                 let dict: &Py<PyDict> = unsafe { dict_ref.downcast_unchecked_ref() };
 
-                // Get callable for error messages
-                // Stack: [callable, self_or_null, args_tuple, kwargs_dict]
-                let callable = self.nth_value(idx + 2);
-                let func_str = Self::object_function_str(callable, vm);
-
-                // Fast path: source is an exact dict (not a subclass, which may
-                // override `keys`/`__getitem__`). Iterate its entries natively
-                // instead of going through the mapping protocol, mirroring
-                // CPython's `PyDict_Merge` fast path for `PyDict_Check(other)`.
-                let source = if source.class().is(vm.ctx.types.dict_type) {
-                    let src_dict = source
-                        .downcast_ref::<PyDict>()
-                        .expect("exact dict must have a PyDict payload");
-                    // Preserve stored hashes: copying an exact dict must not call
-                    // a key's __hash__ before the callable receives its arguments.
-                    // Snapshot under one lock before any equality callbacks run.
-                    for (key, value, hash) in src_dict.items_with_hashes_vec() {
-                        if dict.contains_known_hash(&key, hash, vm)? {
-                            let key_str = key.str(vm)?;
-                            return Err(vm.new_type_error(format!(
-                                "{} got multiple values for keyword argument '{}'",
-                                func_str,
-                                key_str.as_wtf8()
-                            )));
+                let merge_result = (|| -> PyResult<Option<PyObjectRef>> {
+                    // Fast path: source is an exact dict (not a subclass, which may
+                    // override `keys`/`__getitem__`). Iterate its entries natively
+                    // instead of going through the mapping protocol, mirroring
+                    // CPython's `PyDict_Merge` fast path for `PyDict_Check(other)`.
+                    if source.class().is(vm.ctx.types.dict_type) {
+                        let src_dict = source
+                            .downcast_ref::<PyDict>()
+                            .expect("exact dict must have a PyDict payload");
+                        // Preserve stored hashes: copying an exact dict must not call
+                        // a key's __hash__ before the callable receives its arguments.
+                        // Snapshot under one lock before any equality callbacks run.
+                        for (key, value, hash) in src_dict.items_with_hashes_vec() {
+                            if dict.contains_known_hash(&key, hash, vm)? {
+                                return Ok(Some(key));
+                            }
+                            dict.set_item_known_hash(&key, hash, value, vm)?;
                         }
-                        dict.set_item_known_hash(&key, hash, value, vm)?;
+                        return Ok(None);
                     }
-                    return Ok(None);
-                } else {
-                    source
-                };
 
-                // Check if source is a mapping
-                if vm
-                    .get_method(source.clone(), vm.ctx.intern_str("keys"))
-                    .is_none()
-                {
-                    return Err(vm.new_type_error(format!(
-                        "{} argument after ** must be a mapping, not {}",
-                        func_str,
-                        source.class().name()
-                    )));
-                }
+                    // Merge keys, checking for duplicates.
+                    let keys = source.mapping_unchecked().keys(vm)?;
+                    let keys_iter = PyIter::try_from_object(vm, keys)?;
+                    while let PyIterReturn::Return(key) = keys_iter.next(vm)? {
+                        let hash = PyDict::hash_or_unhashable(&*key, vm)?;
+                        if dict.entries.contains(vm, &*key, hash)? {
+                            return Ok(Some(key));
+                        }
+                        let value = source.get_item(&*key, vm)?;
+                        dict.set_item(&*key, value, vm)?;
+                    }
+                    Ok(None)
+                })();
 
-                // Merge keys, checking for duplicates
-                let keys_iter = vm.call_method(&source, "keys", ())?;
-                for key in keys_iter.try_to_value::<Vec<PyObjectRef>>(vm)? {
-                    if dict.contains_key(&*key, vm) {
+                match merge_result {
+                    Ok(Some(key)) => {
+                        // Stack: [callable, self_or_null, args_tuple, kwargs_dict]
+                        let callable = self.nth_value(idx + 2);
+                        let func_str = Self::object_function_str(callable, vm);
                         let key_str = key.str(vm)?;
-                        return Err(vm.new_type_error(format!(
+                        Err(vm.new_type_error(format!(
                             "{} got multiple values for keyword argument '{}'",
                             func_str,
                             key_str.as_wtf8()
-                        )));
+                        )))
                     }
-                    let value = vm.call_method(&source, "__getitem__", (key.clone(),))?;
-                    dict.set_item(&*key, value, vm)?;
+                    Ok(None) => Ok(None),
+                    Err(error) if error.fast_isinstance(vm.ctx.exceptions.attribute_error) => {
+                        // Only missing keys make this a non-mapping. Preserve callback
+                        // errors, chaining a failed recheck as _PyEval_FormatKwargsError does.
+                        match vm.get_attribute_opt(&source, identifier!(vm, keys)) {
+                            Ok(None) => Err(vm.new_type_error(format!(
+                                "Value after ** must be a mapping, not {}",
+                                source.class().fully_qualified_name(vm)?
+                            ))),
+                            Ok(Some(_)) => Err(error),
+                            Err(lookup_error) => {
+                                lookup_error.set_context(Some(error));
+                                Err(lookup_error)
+                            }
+                        }
+                    }
+                    Err(error) => Err(error),
                 }
-                Ok(None)
             }
             Instruction::EndAsyncFor => {
                 // Pops (awaitable, exc) from stack.
@@ -9047,9 +9053,7 @@ impl ExecutingFrame<'_> {
         if let Some(kw_obj) = kwargs_or_null {
             // Stack: [callable, self_or_null, args_tuple]
             let callable = self.nth_value(2);
-            let func_str = Self::object_function_str(callable, vm);
-
-            Self::iterate_mapping_keys(vm, &kw_obj, &func_str, |key, value| {
+            Self::iterate_mapping_keys(vm, &kw_obj, callable, |key, value| {
                 // `PyStr`, not `PyUtf8Str`: CPython only checks that the key is a
                 // `str`, not that it is valid UTF-8, so surrogate keys are accepted.
                 key.downcast_ref::<PyStr>()
@@ -9066,9 +9070,6 @@ impl ExecutingFrame<'_> {
             tuple.as_slice().to_vec()
         } else {
             // Single *arg passed directly; convert to sequence at runtime.
-            // Stack: [callable, self_or_null]
-            let callable = self.nth_value(1);
-            let func_str = Self::object_function_str(callable, vm);
             let not_iterable = args_obj.class().slots().iter.load().is_none()
                 && args_obj
                     .get_class_attr(vm.ctx.intern_str("__getitem__"))
@@ -9076,8 +9077,7 @@ impl ExecutingFrame<'_> {
             args_obj.try_to_value::<Vec<PyObjectRef>>(vm).map_err(|e| {
                 if not_iterable && e.class().is(vm.ctx.exceptions.type_error) {
                     vm.new_type_error(format!(
-                        "{} argument after * must be an iterable, not {}",
-                        func_str,
+                        "Value after * must be an iterable, not {}",
                         args_obj.class().slot_name()
                     ))
                 } else {
@@ -9119,7 +9119,7 @@ impl ExecutingFrame<'_> {
     fn iterate_mapping_keys<F>(
         vm: &VirtualMachine,
         mapping: &PyObject,
-        func_str: &Wtf8,
+        callable: &PyObject,
         mut key_handler: F,
     ) -> PyResult<()>
     where
@@ -9143,12 +9143,17 @@ impl ExecutingFrame<'_> {
         let Some(keys_method) = vm.get_method(mapping.to_owned(), vm.ctx.intern_str("keys")) else {
             return Err(vm.new_type_error(format!(
                 "{} argument after ** must be a mapping, not {}",
-                func_str,
+                Self::object_function_str(callable, vm),
                 mapping.class().name()
             )));
         };
 
-        let keys = PyIter::try_from_object(vm, keys_method?.call((), vm)?)?;
+        let keys = mapping.mapping_unchecked().collect_method_output(
+            identifier!(vm, keys),
+            keys_method?.call((), vm)?,
+            vm,
+        )?;
+        let keys = PyIter::try_from_object(vm, keys)?;
         while let PyIterReturn::Return(key) = keys.next(vm)? {
             let value = mapping.get_item(&*key, vm)?;
             key_handler(key, value)?;

crates/vm/src/protocol/mapping.rs

@@ -212,21 +212,35 @@ impl PyMapping<'_> {
         vm: &VirtualMachine,
     ) -> PyResult {
         let meth_output = vm.call_method(self.obj, method_name.as_str(), ())?;
-        if meth_output.is(vm.ctx.types.list_type) {
+        self.collect_method_output(method_name, meth_output, vm)
+    }
+
+    pub(crate) fn collect_method_output(
+        self,
+        method_name: &'static PyStrInterned,
+        meth_output: PyObjectRef,
+        vm: &VirtualMachine,
+    ) -> PyResult {
+        if meth_output.class().is(vm.ctx.types.list_type) {
             return Ok(meth_output);
         }
 
-        let iter = meth_output.get_iter(vm).map_err(|_| {
-            vm.new_type_error(format!(
-                "{}.{}() returned a non-iterable (type {})",
-                self.obj.class().slot_name(),
-                method_name.as_str(),
-                meth_output.class().slot_name()
-            ))
-        })?;
+        let iter = match meth_output.get_iter(vm) {
+            Ok(iter) => iter,
+            Err(error) if error.fast_isinstance(vm.ctx.exceptions.type_error) => {
+                return Err(vm.new_type_error(format!(
+                    "{}.{}() must return an iterable, not {}",
+                    self.obj.class().fully_qualified_name(vm)?,
+                    method_name.as_str(),
+                    meth_output.class().fully_qualified_name(vm)?
+                )));
+            }
+            Err(error) => return Err(error),
+        };
 
-        // TODO
-        // PySequence::from(&iter).list(vm).map(|x| x.into())
-        vm.ctx.new_list(iter.try_to_value(vm)?).to_pyresult(vm)
+        drop(meth_output);
+        let result = vm.ctx.new_list(Vec::new());
+        result.extend(iter.into(), vm)?;
+        Ok(result.into())
     }
 }