Handle Python2 intern better

Author: rocky

test/bytecode_2.7/01_extended_arg.pyc

@@ -0,0 +1 @@
+../01_extended_arg.pyc

test/bytecode_2.7/roundtrip_pyc/01_extended_arg.pyc

@@ -142,6 +142,12 @@ def codeType2Portable(code, version_tuple=PYTHON_VERSION_TRIPLE):
             co_lnotab=line_table,
             co_freevars=code.co_freevars,  # not in 1.x
             co_cellvars=code.co_cellvars,  # not in 1.x
+
+            # THINK ABOUT: If collection_order isn't defined, i.e. native code
+            # type, should we try to extract it?
+            collection_order=code.collection_order if hasattr(code, "collection_order") else {},
+            reference_objects=code.reference_objects if hasattr(code, "reference_objects") else set(),
+
         )
     else:
         # 1.0 .. 1.5

xdis/codetype/__init__.py

@@ -17,6 +17,7 @@
 import types
 from copy import deepcopy
 from types import CodeType
+from typing import Any, Dict, Set, Tuple, Union
 
 from xdis.codetype.code15 import Code15, Code15FieldTypes
 from xdis.version_info import PYTHON_VERSION_TRIPLE, version_tuple_to_str
@@ -55,20 +56,22 @@ class Code2(Code15):
 
     def __init__(
         self,
-        co_argcount,
-        co_nlocals,
-        co_stacksize,
+        co_argcount: int,
+        co_nlocals: int,
+        co_stacksize: int,
         co_flags,
         co_code,
         co_consts,
         co_names,
         co_varnames,
         co_filename,
-        co_name,
-        co_firstlineno,
+        co_name: str,
+        co_firstlineno: int,
         co_lnotab,
         co_freevars,
         co_cellvars,
+        collection_order: Dict[Union[set, frozenset, dict], Tuple[Any]] = {},
+        reference_objects: Set[Any] = set(),
     ) -> None:
         # Keyword argument parameters in the call below is more robust.
         # Since things change around, robustness is good.
@@ -89,6 +92,19 @@ def __init__(
         self.co_freevars = co_freevars
         self.co_cellvars = co_cellvars
         self.fieldtypes = Code2FieldTypes
+
+        # The following fields are mostly useful in marshaling a code object.
+        # Keeping marshal order exactly the same is useful in round-trip marshal
+        # testing; but it may also have other benefits.
+
+        # By saving the order in sets, frozensets, and dictionary keys,
+        # these collections can be written in the same order that appeared
+        # in unmarshalling (if that's how the code object was created).
+        self.collection_order = collection_order
+
+        # Keeping track of which objects were referenced, allows
+        self.reference_objects = reference_objects
+
         if type(self) is Code2:
             self.check()
         return

xdis/codetype/code20.py

@@ -46,10 +46,10 @@ class Code3(Code2):
 
     def __init__(
         self,
-        co_argcount,
-        co_kwonlyargcount,
-        co_nlocals,
-        co_stacksize,
+        co_argcount: int,
+        co_kwonlyargcount: int,
+        co_nlocals: int,
+        co_stacksize: int,
         co_flags,
         co_code,
         co_consts,
@@ -81,23 +81,12 @@ def __init__(
             co_lnotab=co_lnotab,
             co_freevars=co_freevars,
             co_cellvars=co_cellvars,
+            collection_order=collection_order,
+            reference_objects=reference_objects,
         )
         self.co_kwonlyargcount = co_kwonlyargcount
         self.fieldtypes = Code3FieldTypes
 
-        # The following fields are mostly useful in marshaling a code object.
-        # Keeping marshal order exactly the same is useful in round-trip marshal
-        # testing; but it may also have other benefits.
-
-        # By saving the order in sets, frozensets, and dictionary keys,
-        # these collections can be written in the same order that appeared
-        # in unmarshalling (if that's how the code object was created).
-        self.collection_order = collection_order
-
-        # Keeping track of which objects were referenced, allows
-        self.reference_objects = reference_objects
-
-
         if type(self) is Code3:
             self.check()
         return

xdis/codetype/code30.py

@@ -299,7 +299,7 @@ def dump_collection(self, type_code: str, bag: Union[frozenset, set, dict]) -> N
         for each in collection:
             self.dump(each)
 
-    def dump_complex(self, x) -> None:
+    def dump_complex(self, x, _) -> None:
         write = self._write
         write(TYPE_COMPLEX)
         s = repr(x.real)
@@ -551,16 +551,20 @@ def dump_stopiter(self, x) -> None:
 
     dispatch[type(StopIteration)] = dump_stopiter
 
-    def dump_string(self, x, flag_ref: int = 0) -> None:
+    def dump_string(self, s, flag_ref: int = 0) -> None:
         # Python 3.11 seems to add the object ref flag bit for strings.
-        type_string = (
-            TYPE_STRING
-            if self.python_version < (3, 11)
-            else chr(ord(TYPE_STRING) | flag_ref)
-        )
-        self._write(type_string)
-        self.w_long(len(x))
-        self._write(x)
+        if self.python_version >= (3, 11):
+            type_code = chr(ord(TYPE_STRING) | flag_ref)
+        if (3, 0) <= self.python_version < (3, 11):
+            type_code = TYPE_STRING
+        else:
+            # Python 2.x.
+            # FIXME: save string somewhere if it isn't in string table.
+            type_code = TYPE_INTERNED if s in self.reference_objects else TYPE_STRING
+
+        self._write(type_code)
+        self.w_long(len(s))
+        self._write(s)
 
     dispatch[bytes] = dump_string
     dispatch[bytearray] = dump_string
@@ -585,8 +589,10 @@ def dump_tuple(self, tuple_object: tuple, flag_ref: int = 0) -> None:
     def dump_unicode(self, s, flag_ref: int = 0) -> None:
         if self.python_version < (2, 0):
             type_code = TYPE_STRING
-        elif (2, 0) <= self.python_version <= (3, 0):
-            type_code = TYPE_INTERNED
+        elif (2, 0) <= self.python_version < (3, 0):
+            # FIXME: probably need to save string somewhere
+            # if it isn't in string table.
+            type_code = TYPE_INTERNED if s in self.reference_objects else TYPE_STRING
         else:
             type_code = TYPE_UNICODE
 

xdis/marsh.py

@@ -398,7 +398,6 @@ def t_short_ASCII_interned(self, save_ref, bytes_for_s: bool = False):
         self.intern_strings.append(interned)
         return self.r_ref(interned, save_ref)
 
-    # Since Python 3.4
     def t_interned(self, save_ref, bytes_for_s: bool = False):
         strsize = unpack("<i", self.fp.read(4))[0]
         interned = compat_str(self.fp.read(strsize))