cosmopolitan/third_party/python/freeze.c

/*-*- mode:c;indent-tabs-mode:nil;c-basic-offset:4;tab-width:8;coding:utf-8 -*-│
│vi: set net ft=c ts=4 sts=4 sw=4 fenc=utf-8                                :vi│
╞══════════════════════════════════════════════════════════════════════════════╡
│ Python 3                                                                     │
│ https://docs.python.org/3/license.html                                       │
╚─────────────────────────────────────────────────────────────────────────────*/
#include "libc/assert.h"
#include "libc/mem/mem.h"
#include "third_party/python/Include/bytesobject.h"
#include "third_party/python/Include/compile.h"
#include "third_party/python/Include/fileutils.h"
#include "third_party/python/Include/import.h"
#include "third_party/python/Include/marshal.h"
#include "third_party/python/Include/pydebug.h"
#include "third_party/python/Include/pylifecycle.h"
#include "third_party/python/Include/pymacro.h"
#include "third_party/python/Include/pythonrun.h"
/* clang-format off */

#define HEADER "\
/*-*- mode:c;indent-tabs-mode:nil;c-basic-offset:4;tab-width:8;coding:utf-8 -*-│\n\
│vi: set net ft=c ts=4 sts=4 sw=4 fenc=utf-8                                :vi│\n\
╞══════════════════════════════════════════════════════════════════════════════╡\n\
│ Python 3                                                                     │\n\
│ https://docs.python.org/3/license.html                                       │\n\
╚─────────────────────────────────────────────────────────────────────────────*/\n\
/* clang-format off */\n\
\n\
/*\n\
 * Auto-generated by\n\
 * o/third_party/python/freeze \\\n\
 *   %s \\\n\
 *   %s\n\
 */\n\
\n\
"

int
main(int argc, char *argv[])
{
    char *inpath, *outpath, *code_name;
    FILE *infile = NULL, *outfile = NULL;
    struct _Py_stat_struct status;
    size_t text_size, data_size, n;
    char *text = NULL;
    unsigned char *data;
    PyObject *code = NULL, *marshalled = NULL;
    int is_bootstrap = 1;

    if (argc == 2 && !strcmp(argv[1], "-n")) return 0;

    if (argc != 3) {
        fprintf(stderr, "need to specify input and output paths\n");
        return 2;
    }
    inpath = argv[1];
    outpath = argv[2];
    infile = fopen(inpath, "rb");
    if (infile == NULL) {
        fprintf(stderr, "cannot open '%s' for reading\n", inpath);
        goto error;
    }
    if (_Py_fstat_noraise(fileno(infile), &status)) {
        fprintf(stderr, "cannot fstat '%s'\n", inpath);
        goto error;
    }
    text_size = (size_t)status.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';

    Py_NoUserSiteDirectory++;
    Py_NoSiteFlag++;
    Py_IgnoreEnvironmentFlag++;
    Py_FrozenFlag++;

    Py_SetProgramName(L"./_freeze_importlib");
    /* Don't install importlib, since it could execute outdated bytecode. */
    _Py_InitializeEx_Private(1, 0);

    if (strstr(inpath, "_external") != NULL) {
        is_bootstrap = 0;
    }

    code_name = is_bootstrap ?
        "<frozen importlib._bootstrap>" :
        "<frozen importlib._bootstrap_external>";
    code = Py_CompileStringExFlags(text, code_name, Py_file_input, NULL, 2);
    if (code == NULL)
        goto error;
    free(text);
    text = NULL;

    marshalled = PyMarshal_WriteObjectToString(code, Py_MARSHAL_VERSION);
    Py_CLEAR(code);
    if (marshalled == NULL)
        goto error;

    assert(PyBytes_CheckExact(marshalled));
    data = (unsigned char *) PyBytes_AS_STRING(marshalled);
    data_size = PyBytes_GET_SIZE(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 */
    outfile = fopen(outpath, "w");
    if (outfile == NULL) {
        fprintf(stderr, "cannot open '%s' for writing\n", outpath);
        goto error;
    }
    fprintf(outfile, HEADER, argv[1], argv[2]);
    if (is_bootstrap)
        fprintf(outfile, "const unsigned char _Py_M__importlib[] = {\n");
    else
        fprintf(outfile,
                "const unsigned char _Py_M__importlib_external[] = {\n");
    for (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++) {
            fprintf(outfile, "%d,", (unsigned int) data[i]);
        }
        fprintf(outfile, "\n");
    }
    fprintf(outfile, "};\n");

    Py_CLEAR(marshalled);

    Py_Finalize();
    if (outfile) {
        if (ferror(outfile)) {
            fprintf(stderr, "error when writing to '%s'\n", outpath);
            goto error;
        }
        fclose(outfile);
    }
    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;
}