Match collection ordering in unmarshal to marshal

Author: rocky

test/bytecode_3.3/roundtrip_pyc/06_frozenset.pyc

@@ -0,0 +1 @@
+../06_frozenset.pyc

xdis/codetype/__init__.py

@@ -16,7 +16,7 @@
 
 import types
 from collections import namedtuple
-from typing import Optional
+from typing import Any, Dict, Optional, Tuple, Union
 
 from xdis.codetype.base import CodeBase
 from xdis.codetype.code13 import Code13
@@ -60,6 +60,10 @@ def codeType2Portable(code, version_tuple=PYTHON_VERSION_TRIPLE):
                 line_table,
                 code.co_freevars,
                 code.co_cellvars,
+
+                # THINK ABOUT: If collection_order isn't defined, i.e. native code
+                # type, should we try to extract it?
+                code.collection_order if hasattr(code, "collection_order") else {}
             )
         elif version_tuple < (3, 10):
             return Code38(
@@ -203,7 +207,7 @@ def portableCodeType(version_tuple=PYTHON_VERSION_TRIPLE):
 # In contrast to Code3, Code2, etc. you can use CodeTypeUnint for building
 # an incomplete code type, which might be converted to another code type
 # later.
-CodeTypeUnionFields = tuple(Code311FieldNames.split())
+CodeTypeUnionFields = tuple(Code311FieldNames.split() + ["collection_order"])
 CodeTypeUnion = namedtuple("CodeTypeUnion", CodeTypeUnionFields)
 
 
@@ -230,6 +234,7 @@ def to_portable(
     co_cellvars: tuple[None]=(None,),  # 2.0+
     co_exceptiontable=None,  # 3.11+
     version_triple=PYTHON_VERSION_TRIPLE,
+    collection_order: Dict[Union[set, frozenset, dict], Tuple[Any]] = {},
 ):
     code = CodeTypeUnion(
         co_argcount=co_argcount,
@@ -250,6 +255,7 @@ def to_portable(
         co_freevars=co_freevars,
         co_cellvars=co_cellvars,
         co_exceptiontable=co_exceptiontable,
+        collection_order=collection_order,
     )
     return codeType2Portable(code, version_triple)
 

xdis/codetype/code20.py

@@ -72,7 +72,7 @@ def __init__(
     ) -> None:
         # Keyword argument parameters in the call below is more robust.
         # Since things change around, robustness is good.
-        super(Code2, self).__init__(
+        super().__init__(
             co_argcount=co_argcount,
             co_nlocals=co_nlocals,
             co_stacksize=co_stacksize,
@@ -89,7 +89,7 @@ def __init__(
         self.co_freevars = co_freevars
         self.co_cellvars = co_cellvars
         self.fieldtypes = Code2FieldTypes
-        if type(self) == Code2:
+        if type(self) is Code2:
             self.check()
         return
 

xdis/codetype/code30.py

@@ -17,7 +17,7 @@
 import types
 from copy import deepcopy
 from types import CodeType
-from typing import Dict, List, Union
+from typing import Any, Dict, Tuple, Union
 
 from xdis.codetype.code20 import Code2, Code2FieldTypes
 from xdis.version_info import PYTHON_VERSION_TRIPLE, version_tuple_to_str
@@ -61,7 +61,7 @@ def __init__(
         co_lnotab,
         co_freevars,
         co_cellvars,
-        collection_order: Dict[Union[set, frozenset, dict], List[str]] = {}
+        collection_order: Dict[Union[set, frozenset, dict], Tuple[Any]] = {}
     ) -> None:
         # Keyword argument parameters in the call below is more robust.
         # Since things change around, robustness is good.

xdis/codetype/code310.py

@@ -17,6 +17,7 @@
 import struct
 import types
 from copy import deepcopy
+from typing import Any, Dict, Tuple, Union
 
 from xdis.codetype.code38 import Code38
 from xdis.cross_types import UnicodeForPython3
@@ -96,6 +97,7 @@ def __init__(
         co_linetable,
         co_freevars,
         co_cellvars,
+        collection_order: Dict[Union[set, frozenset, dict], Tuple[Any]] = {}
     ) -> None:
         # Keyword argument parameters in the call below is more robust.
         # Since things change around, robustness is good.
@@ -116,6 +118,11 @@ def __init__(
         self.co_stacksize = co_stacksize
         self.co_varnames = co_varnames
         self.fieldtypes = Code310FieldTypes
+
+        # It is helpful to save the order in sets, frozensets and dictionary keys,
+        # so that on writing a bytecode file we can duplicate this order.
+        self.collection_order = collection_order
+
         if type(self) is Code310:
             self.check()
 

xdis/codetype/code38.py

@@ -16,7 +16,7 @@
 
 import types
 from copy import deepcopy
-from typing import Dict, List, Union
+from typing import Any, Dict, Tuple, Union
 
 from xdis.codetype.code30 import Code3, Code3FieldTypes
 from xdis.version_info import PYTHON_VERSION_TRIPLE, version_tuple_to_str
@@ -84,7 +84,7 @@ def __init__(
         co_lnotab,
         co_freevars,
         co_cellvars,
-        collection_order: Dict[Union[set, frozenset, dict], List[str]] = {}
+        collection_order: Dict[Union[set, frozenset, dict], Tuple[Any]] = {}
     ) -> None:
         # Keyword argument parameters in the call below is more robust.
         # Since things change around, robustness is good.

xdis/marsh.py

@@ -28,7 +28,7 @@
 import types
 from sys import intern
 from types import CodeType, EllipsisType
-from typing import Optional
+from typing import Any, Optional, Set, Union
 
 from xdis.codetype import Code2, Code3, Code15
 from xdis.unmarshal import (
@@ -88,11 +88,16 @@ class _Marshaller:
     dispatch = {}
 
     def __init__(
-        self, writefunc, python_version: tuple, is_pypy: Optional[bool] = None
+        self,
+        writefunc,
+        python_version: tuple,
+        is_pypy: Optional[bool] = None,
+        collection_order={},
     ) -> None:
         self._write = writefunc
         self.python_version = python_version
         self.is_pypy = is_pypy
+        self.collection_order = collection_order
 
     def dump(self, x) -> None:
         if (
@@ -417,6 +422,7 @@ def dump_code3(self, x) -> None:
 
     dispatch[Code3] = dump_code3
 
+    # FIXME: this is wrong.
     try:
         if PYTHON3:
             dispatch[types.CodeType] = dump_code3
@@ -425,22 +431,32 @@ def dump_code3(self, x) -> None:
     except NameError:
         pass
 
-    def dump_set(self, x) -> None:
-        self._write(TYPE_SET)
-        self.w_long(len(x))
-        for each in x:
+    def dump_collection(self, type_code: str, bag: Union[frozenset, set, dict]) -> None:
+        """
+        Save marshalled version of frozenset fs.
+        Use self.collection_order, to ensure that the order
+        or set elements that may have appeared from unmarshalling the appears
+        the same way. This helps roundtrip checking, among possibly other things.
+        """
+        self._write(type_code)
+        self.w_long(len(bag))
+        collection = self.collection_order.get(bag, bag)
+        for each in collection:
             self.dump(each)
 
-    try:
-        dispatch[set] = dump_set
-    except NameError:
-        pass
+    def dump_set(self, s: Set[Any]) -> None:
+        """
+        Save marshalled version of set s.
+        """
+        self.dump_collection(TYPE_SET, s)
 
-    def dump_frozenset(self, x) -> None:
-        self._write(TYPE_FROZENSET)
-        self.w_long(len(x))
-        for each in x:
-            self.dump(each)
+    dispatch[set] = dump_set
+
+    def dump_frozenset(self, fs: frozenset) -> None:
+        """
+        Save marshalled version of frozenset fs.
+        """
+        self.dump_collection(TYPE_FROZENSET, fs)
 
     try:
         dispatch[frozenset] = dump_frozenset
@@ -1103,7 +1119,13 @@ def dumps(
     is_pypy: Optional[bool] = None,
 ) -> bytes | str:
     buffer = []
-    m = _Marshaller(buffer.append, python_version=python_version, is_pypy=is_pypy)
+    collection_order = x.collection_order if hasattr(x, "collection_order") else {}
+    m = _Marshaller(
+        buffer.append,
+        python_version=python_version,
+        is_pypy=is_pypy,
+        collection_order=collection_order,
+    )
     m.dump(x)
     if python_version:
         is_python3 = python_version >= (3, 0)

xdis/unmarshal.py

@@ -30,7 +30,7 @@
 import sys
 from struct import unpack
 from types import EllipsisType
-from typing import Union
+from typing import Any, Dict, Tuple, Union
 
 from xdis.codetype import to_portable
 from xdis.cross_types import LongTypeForPython3, UnicodeForPython3
@@ -162,6 +162,10 @@ def __init__(self, fp, magic_int, bytes_for_s, code_objects={}) -> None:
         # objects starts.
         self.code_to_file_offsets = {}
 
+        # It is helpful to save the order in sets, frozensets and dictionary keys,
+        # so that on writing a bytecode file we can duplicate this order.
+        self.collection_order: Dict[Union[set, frozenset, dict], Tuple[Any]] = {}
+
         self.bytes_for_s = bytes_for_s
         version = magic_int2tuple(self.magic_int)
         if version >= (3, 4):
@@ -441,11 +445,14 @@ def t_list(self, save_ref, bytes_for_s: bool = False):
 
     def t_frozenset(self, save_ref, bytes_for_s: bool = False):
         setsize = unpack("<i", self.fp.read(4))[0]
-        ret, i = self.r_ref_reserve(tuple(), save_ref)
+        collection, i = self.r_ref_reserve([], save_ref)
         while setsize > 0:
-            ret += (self.r_object(bytes_for_s=bytes_for_s),)
+            collection.append(self.r_object(bytes_for_s=bytes_for_s))
             setsize -= 1
-        return self.r_ref_insert(frozenset(ret), i)
+        final_frozenset = frozenset(collection)
+        # Note the order of the frozenset elements.
+        self.collection_order[final_frozenset] = tuple(collection)
+        return self.r_ref_insert(final_frozenset, i)
 
     def t_set(self, save_ref, bytes_for_s: bool = False):
         setsize = unpack("<i", self.fp.read(4))[0]
@@ -658,6 +665,7 @@ def t_code(self, save_ref, bytes_for_s: bool = False):
             co_cellvars=co_cellvars,
             co_exceptiontable=co_exceptiontable,
             version_triple=self.version_tuple,
+            collection_order=self.collection_order,
         )
 
         self.code_to_file_offsets[code] = (code_offset_in_file, co_code_offset_in_file)