Small tweaks

Author: rocky

xdis/load.py

@@ -380,7 +380,12 @@ def load_module_from_file_object(
 
 
 def write_bytecode_file(
-    bytecode_path, code_obj, magic_int, compilation_ts=None, filesize: int = 0
+    bytecode_path,
+    code_obj,
+    magic_int,
+    compilation_ts=None,
+    filesize: int = 0,
+    allow_native: bool = True,
 ) -> None:
     """Write bytecode file _bytecode_path_, with code for having Python
     magic_int (i.e. bytecode associated with some version of Python)
@@ -407,15 +412,19 @@ def write_bytecode_file(
     if version_tuple >= (3, 3):
         # In Python 3.3+, these 4 bytes are the size of the source code_obj file (mod 2^32)
         fp.write(pack("<I", filesize))
-    if isinstance(code_obj, types.CodeType):
+    if allow_native and isinstance(code_obj, types.CodeType):
         fp.write(marshal.dumps(code_obj))
     else:
         code_sequence = xdis.marsh.dumps(code_obj, python_version=version_tuple)
         if isinstance(code_sequence, str):
             # Python 1.x uses code strings, not bytes. To get this into bytes needed by
             # fp.write, encode the string using 'latin-1' and 'unicode_escape' to convert escape sequences
             # into the raw byte values. 'latin-1' is a single-byte encoding that works well for this.
-            code_bytes = code_sequence.encode('latin-1').decode('unicode_escape').encode('latin-1')
+            code_bytes = (
+                code_sequence.encode("latin-1")
+                .decode("unicode_escape")
+                .encode("latin-1")
+            )
         else:
             code_bytes = code_sequence
         fp.write(code_bytes)
@@ -425,8 +434,8 @@ def write_bytecode_file(
 if __name__ == "__main__":
     co = load_file(__file__)
     obj_path = check_object_path(__file__)
-    version, timestamp, magic_int, co2, pypy, source_size, sip_hash, file_offsets = load_module(
-        obj_path
+    version, timestamp, magic_int, co2, pypy, source_size, sip_hash, file_offsets = (
+        load_module(obj_path)
     )
     print("version", version, "magic int", magic_int, "is_pypy", pypy)
     if timestamp is not None:

xdis/roundtrip_pyc.py

@@ -40,7 +40,7 @@
 import sys
 import tempfile
 import types
-from typing import List, Union
+from typing import Union
 
 from xdis.codetype.base import CodeBase
 from xdis.load import load_module_from_file_object, write_bytecode_file
@@ -231,6 +231,7 @@ def roundtrip_pyc(input_path: str, unlink_on_success: bool, verbose: bool) -> in
             orig_magic_int,
             compilation_ts=orig_timestamp,
             filesize=orig_source_size or 0,
+            allow_native=False,
         )
     except TypeError:
         # Older/newer signatures might name the timestamp param differently; try without names
@@ -261,7 +262,7 @@ def roundtrip_pyc(input_path: str, unlink_on_success: bool, verbose: bool) -> in
     except Exception as e:
         print(f"WARNING: could not do raw byte comparison: {e}", file=sys.stderr)
 
-    print("Original file:", input_path)
+    print("Input file:", input_path)
     if verbose:
         print("Rewritten file:", tf_name)
         print("Raw-bytes identical:", same_bytes)