bpo-46429: Merge all deepfrozen files into one (GH-30572)

Tools/scripts/deepfreeze.py

@@ -10,10 +10,9 @@ import collections
 import contextlib
 import os
 import re
-import sys
 import time
 import types
-from typing import Dict, FrozenSet, Tuple, TextIO
+from typing import Dict, FrozenSet, TextIO, Tuple
 
 import umarshal
 
@@ -104,10 +103,10 @@ def removesuffix(base: str, suffix: str) -> str:
 
 class Printer:
 
-    def __init__(self, file: TextIO):
+    def __init__(self, file: TextIO) -> None:
         self.level = 0
         self.file = file
-        self.cache: Dict[Tuple[type, object], str] = {}
+        self.cache: Dict[tuple[type, object, str], str] = {}
         self.hits, self.misses = 0, 0
         self.patchups: list[str] = []
         self.write('#include "Python.h"')
@@ -349,6 +348,15 @@ class Printer:
         self.write("// TODO: The above tuple should be a frozenset")
         return ret
 
+    def generate_file(self, module: str, code: object)-> None:
+        module = module.replace(".", "_")
+        self.generate(f"{module}_toplevel", code)
+        with self.block(f"static void {module}_do_patchups(void)"):
+            for p in self.patchups:
+                self.write(p)
+        self.patchups.clear()
+        self.write(EPILOGUE.replace("%%NAME%%", module))
+
     def generate(self, name: str, obj: object) -> str:
         # Use repr() in the key to distinguish -0.0 from +0.0
         key = (type(obj), obj, repr(obj))
@@ -357,7 +365,7 @@ class Printer:
             # print(f"Cache hit {key!r:.40}: {self.cache[key]!r:.40}")
             return self.cache[key]
         self.misses += 1
-        if isinstance(obj, types.CodeType) or isinstance(obj, umarshal.Code):
+        if isinstance(obj, (types.CodeType, umarshal.Code)) :
             val = self.generate_code(name, obj)
         elif isinstance(obj, tuple):
             val = self.generate_tuple(name, obj)
@@ -393,8 +401,8 @@ EPILOGUE = """
 PyObject *
 _Py_get_%%NAME%%_toplevel(void)
 {
-    do_patchups();
-    return (PyObject *) &toplevel;
+    %%NAME%%_do_patchups();
+    return (PyObject *) &%%NAME%%_toplevel;
 }
 """
 
@@ -419,29 +427,25 @@ def decode_frozen_data(source: str) -> types.CodeType:
     return umarshal.loads(data)
 
 
-def generate(source: str, filename: str, modname: str, file: TextIO) -> None:
-    if is_frozen_header(source):
-        code = decode_frozen_data(source)
-    else:
-        code = compile(source, filename, "exec")
-    printer = Printer(file)
-    printer.generate("toplevel", code)
-    printer.write("")
-    with printer.block("static void do_patchups(void)"):
-        for p in printer.patchups:
-            printer.write(p)
-    here = os.path.dirname(__file__)
-    printer.write(EPILOGUE.replace("%%NAME%%", modname.replace(".", "_")))
+def generate(args: list[str], output: TextIO) -> None:
+    printer = Printer(output)
+    for arg in args:
+        file, modname = arg.rsplit(':', 1)
+        with open(file, "r", encoding="utf8") as fd:
+            source = fd.read()
+            if is_frozen_header(source):
+                code = decode_frozen_data(source)
+            else:
+                code = compile(fd.read(), f"<frozen {modname}>", "exec")
+            printer.generate_file(modname, code)
     if verbose:
         print(f"Cache hits: {printer.hits}, misses: {printer.misses}")
 
 
 parser = argparse.ArgumentParser()
-parser.add_argument("-m", "--module", help="Defaults to basename(file)")
-parser.add_argument("-o", "--output", help="Defaults to MODULE.c")
+parser.add_argument("-o", "--output", help="Defaults to deepfreeze.c", default="deepfreeze.c")
 parser.add_argument("-v", "--verbose", action="store_true", help="Print diagnostics")
-parser.add_argument("file", help="Input file (required)")
-
+parser.add_argument('args', nargs="+", help="Input file and module name (required) in file:modname format")
 
 @contextlib.contextmanager
 def report_time(label: str):
@@ -458,13 +462,10 @@ def main() -> None:
     global verbose
     args = parser.parse_args()
     verbose = args.verbose
-    with open(args.file, encoding="utf-8") as f:
-        source = f.read()
-    modname = args.module or removesuffix(os.path.basename(args.file), ".py")
-    output = args.output or modname + ".c"
+    output = args.output
     with open(output, "w", encoding="utf-8") as file:
         with report_time("generate"):
-            generate(source, f"<frozen {modname}>", modname, file)
+            generate(args.args, file)
     if verbose:
         print(f"Wrote {os.path.getsize(output)} bytes to {output}")