bpo-45019: Add a tool to generate list of modules to include for frozen modules (gh-27980)

Frozen modules must be added to several files in order to work properly. Before this change this had to be done manually. Here we add a tool to generate the relevant lines in those files instead. This helps us avoid mistakes and omissions.

https://bugs.python.org/issue45019

1 files changed, 149 insertions, 90 deletions

diff --git a/Programs/_freeze_importlib.c b/Programs/_freeze_module.c
index 2e4ccbb..7e9f02a 100644
--- a/Programs/_freeze_importlib.c
+++ b/Programs/_freeze_module.c
@@ -1,5 +1,10 @@
 /* This is built as a stand-alone executable by the Makefile, and helps turn
-   Lib/importlib/_bootstrap.py into a frozen module in Python/importlib.h
+   modules into frozen modules (like Lib/importlib/_bootstrap.py
+   into Python/importlib.h).
+
+   This is used directly by Tools/scripts/freeze_modules.py, and indirectly by "make regen-frozen".
+
+   See Python/frozen.c for more info.
 */
 
 #include <Python.h>
@@ -28,54 +33,11 @@ const struct _frozen *PyImport_FrozenModules;
 #endif
 
 static const char header[] =
-    "/* Auto-generated by Programs/_freeze_importlib.c */";
+    "/* Auto-generated by Programs/_freeze_module.c */";
 
-int
-main(int argc, char *argv[])
+static void
+runtime_init(void)
 {
-    const char *name, *inpath, *outpath;
-    char buf[100];
-    FILE *infile = NULL, *outfile = NULL;
-    struct _Py_stat_struct stat;
-    size_t text_size, data_size, i, n;
-    char *text = NULL;
-    unsigned char *data;
-    PyObject *code = NULL, *marshalled = NULL;
-
-    PyImport_FrozenModules = _PyImport_FrozenModules;
-
-    if (argc != 4) {
-        fprintf(stderr, "need to specify the name, input and output paths\n");
-        return 2;
-    }
-    name = argv[1];
-    inpath = argv[2];
-    outpath = argv[3];
-    infile = fopen(inpath, "rb");
-    if (infile == NULL) {
-        fprintf(stderr, "cannot open '%s' for reading\n", inpath);
-        goto error;
-    }
-    if (_Py_fstat_noraise(fileno(infile), &stat)) {
-        fprintf(stderr, "cannot fstat '%s'\n", inpath);
-        goto error;
-    }
-    text_size = (size_t)stat.st_size;
-    text = (char *) malloc(text_size + 1);
-    if (text == NULL) {
-        fprintf(stderr, "could not allocate %ld bytes\n", (long) text_size);
-        goto error;
-    }
-    n = fread(text, 1, text_size, infile);
-    fclose(infile);
-    infile = NULL;
-    if (n < text_size) {
-        fprintf(stderr, "read too short: got %ld instead of %ld bytes\n",
-                (long) n, (long) text_size);
-        goto error;
-    }
-    text[text_size] = '\0';
-
     PyConfig config;
     PyConfig_InitIsolatedConfig(&config);
 
@@ -83,7 +45,7 @@ main(int argc, char *argv[])
 
     PyStatus status;
     status = PyConfig_SetString(&config, &config.program_name,
-                                L"./_freeze_importlib");
+                                L"./_freeze_module");
     if (PyStatus_Exception(status)) {
         PyConfig_Clear(&config);
         Py_ExitStatusException(status);
@@ -98,39 +60,93 @@ main(int argc, char *argv[])
     if (PyStatus_Exception(status)) {
         Py_ExitStatusException(status);
     }
+}
 
-    sprintf(buf, "<frozen %s>", name);
-    code = Py_CompileStringExFlags(text, buf, Py_file_input, NULL, 0);
-    if (code == NULL)
-        goto error;
-    free(text);
-    text = NULL;
+static const char *
+read_text(const char *inpath)
+{
+    FILE *infile = fopen(inpath, "rb");
+    if (infile == NULL) {
+        fprintf(stderr, "cannot open '%s' for reading\n", inpath);
+        return NULL;
+    }
 
-    marshalled = PyMarshal_WriteObjectToString(code, Py_MARSHAL_VERSION);
-    Py_CLEAR(code);
-    if (marshalled == NULL)
-        goto error;
+    struct _Py_stat_struct stat;
+    if (_Py_fstat_noraise(fileno(infile), &stat)) {
+        fprintf(stderr, "cannot fstat '%s'\n", inpath);
+        fclose(infile);
+        return NULL;
+    }
+    size_t text_size = (size_t)stat.st_size;
 
-    assert(PyBytes_CheckExact(marshalled));
-    data = (unsigned char *) PyBytes_AS_STRING(marshalled);
-    data_size = PyBytes_GET_SIZE(marshalled);
+    char *text = (char *) malloc(text_size + 1);
+    if (text == NULL) {
+        fprintf(stderr, "could not allocate %ld bytes\n", (long) text_size);
+        fclose(infile);
+        return NULL;
+    }
+    size_t n = fread(text, 1, text_size, infile);
+    fclose(infile);
 
-    /* Open the file in text mode. The hg checkout should be using the eol extension,
-       which in turn should cause the EOL style match the C library's text mode */
-    outfile = fopen(outpath, "w");
-    if (outfile == NULL) {
-        fprintf(stderr, "cannot open '%s' for writing\n", outpath);
-        goto error;
+    if (n < text_size) {
+        fprintf(stderr, "read too short: got %ld instead of %ld bytes\n",
+                (long) n, (long) text_size);
+        free(text);
+        return NULL;
     }
-    fprintf(outfile, "%s\n", header);
-    for (i = n = 0; name[i] != '\0'; i++) {
-        if (name[i] != '.') {
-            buf[n++] = name[i];
+
+    text[text_size] = '\0';
+    return (const char *)text;
+}
+
+static PyObject *
+compile_and_marshal(const char *name, const char *text)
+{
+    char *filename = (char *) malloc(strlen(name) + 10);
+    sprintf(filename, "<frozen %s>", name);
+    PyObject *code = Py_CompileStringExFlags(text, filename,
+                                             Py_file_input, NULL, 0);
+    free(filename);
+    if (code == NULL) {
+        return NULL;
+    }
+
+    PyObject *marshalled = PyMarshal_WriteObjectToString(code, Py_MARSHAL_VERSION);
+    Py_CLEAR(code);
+    if (marshalled == NULL) {
+        return NULL;
+    }
+    assert(PyBytes_CheckExact(marshalled));
+
+    return marshalled;
+}
+
+static char *
+get_varname(const char *name, const char *prefix)
+{
+    size_t n = strlen(prefix);
+    char *varname = (char *) malloc(strlen(name) + n + 1);
+    (void)strcpy(varname, prefix);
+    for (size_t i = 0; name[i] != '\0'; i++) {
+        if (name[i] == '.') {
+            varname[n++] = '_';
+        }
+        else {
+            varname[n++] = name[i];
         }
     }
-    buf[n] = '\0';
-    fprintf(outfile, "const unsigned char _Py_M__%s[] = {\n", buf);
-    for (n = 0; n < data_size; n += 16) {
+    varname[n] = '\0';
+    return varname;
+}
+
+static void
+write_code(FILE *outfile, PyObject *marshalled, const char *varname)
+{
+    unsigned char *data = (unsigned char *) PyBytes_AS_STRING(marshalled);
+    size_t data_size = PyBytes_GET_SIZE(marshalled);
+
+    fprintf(outfile, "const unsigned char %s[] = {\n", varname);
+    for (size_t n = 0; n < data_size; n += 16) {
         size_t i, end = Py_MIN(n + 16, data_size);
         fprintf(outfile, "    ");
         for (i = n; i < end; i++) {
@@ -139,29 +155,72 @@ main(int argc, char *argv[])
         fprintf(outfile, "\n");
     }
     fprintf(outfile, "};\n");
+}
+
+static int
+write_frozen(const char *outpath, const char *inpath, const char *name,
+             PyObject *marshalled)
+{
+    /* Open the file in text mode. The hg checkout should be using the eol extension,
+       which in turn should cause the EOL style match the C library's text mode */
+    FILE *outfile = fopen(outpath, "w");
+    if (outfile == NULL) {
+        fprintf(stderr, "cannot open '%s' for writing\n", outpath);
+        return -1;
+    }
 
-    Py_CLEAR(marshalled);
+    fprintf(outfile, "%s\n", header);
+    char *arrayname = get_varname(name, "_Py_M__");
+    write_code(outfile, marshalled, arrayname);
+    free(arrayname);
 
-    Py_Finalize();
-    if (outfile) {
-        if (ferror(outfile)) {
-            fprintf(stderr, "error when writing to '%s'\n", outpath);
-            goto error;
-        }
-        fclose(outfile);
+    if (ferror(outfile)) {
+        fprintf(stderr, "error when writing to '%s'\n", outpath);
+        return -1;
     }
+    fclose(outfile);
+    return 0;
+}
+
+int
+main(int argc, char *argv[])
+{
+    const char *name, *inpath, *outpath;
+
+    PyImport_FrozenModules = _PyImport_FrozenModules;
+
+    if (argc != 4) {
+        fprintf(stderr, "need to specify the name, input and output paths\n");
+        return 2;
+    }
+    name = argv[1];
+    inpath = argv[2];
+    outpath = argv[3];
+
+    runtime_init();
+
+    const char *text = read_text(inpath);
+    if (text == NULL) {
+        goto error;
+    }
+
+    PyObject *marshalled = compile_and_marshal(name, text);
+    free((char *)text);
+    if (marshalled == NULL) {
+        goto error;
+    }
+
+    int res = write_frozen(outpath, inpath, name, marshalled);
+    Py_DECREF(marshalled);
+    if (res != 0) {
+        goto error;
+    }
+
+    Py_Finalize();
     return 0;
 
 error:
     PyErr_Print();
     Py_Finalize();
-    if (infile)
-        fclose(infile);
-    if (outfile)
-        fclose(outfile);
-    if (text)
-        free(text);
-    if (marshalled)
-        Py_DECREF(marshalled);
     return 1;
 }