From 37c3a50c7fed020724f1e60c3f6cf96205abb899 Mon Sep 17 00:00:00 2001
From: Hood Chatham <roberthoodchatham@gmail.com>
Date: Mon, 20 Jan 2025 16:26:01 +0100
Subject: [PATCH] Add a new subcommand to workerd for making a Python snapshot

For now this only makes baseline snapshots but I'll extend it to make
package snapshots when I need it.
---
 src/workerd/server/workerd.c++ | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/src/workerd/server/workerd.c++ b/src/workerd/server/workerd.c++
index 1dbb1aa1476..efafd34d555 100644
--- a/src/workerd/server/workerd.c++
+++ b/src/workerd/server/workerd.c++
@@ -703,6 +703,9 @@ class CliMain final: public SchemaFileImpl::ErrorReporter {
           .addSubCommand("test", KJ_BIND_METHOD(*this, getTest), "run unit tests")
           .addSubCommand("pyodide-lock", KJ_BIND_METHOD(*this, getPyodideLock),
               "outputs the package lock file used by Pyodide")
+          .addSubCommand("make-pyodide-baseline-snapshot",
+              KJ_BIND_METHOD(*this, getMakePyodideBaselineSnapshot),
+              "Make a Pyodide baseline memory snapshot")
           .build();
       // TODO(someday):
       // "validate": Loads the config and parses all the code to report errors, but then exits
@@ -899,6 +902,18 @@ class CliMain final: public SchemaFileImpl::ErrorReporter {
         .build();
   }
 
+  kj::MainFunc getMakePyodideBaselineSnapshot() {
+    server->allowExperimental();
+    server->setPythonCreateBaselineSnapshot();
+    auto builder =
+        kj::MainBuilder(context, getVersionString(), "Make a Pyodide baseline memory snapshot", "");
+    setPyodideDiskCacheDir(".");
+    return builder.expectArg("<python-version>", CLI_METHOD(parsePythonCompatFlag))
+        .expectArg("<output-directory>", CLI_METHOD(setPackageDiskCacheDir))
+        .callAfterParsing(CLI_METHOD(test))
+        .build();
+  }
+
   void addImportPath(kj::StringPtr pathStr) {
     auto path = fs->getCurrentPath().evalNative(pathStr);
     if (fs->getRoot().tryOpenSubdir(path) != kj::none) {
@@ -1023,6 +1038,24 @@ class CliMain final: public SchemaFileImpl::ErrorReporter {
     server->setPyodideDiskCacheRoot(kj::mv(dir));
   }
 
+  void parsePythonCompatFlag(kj::StringPtr compatFlagStr) {
+    auto builder = kj::heap<capnp::MallocMessageBuilder>();
+    auto configBuilder = builder->initRoot<config::Config>();
+    auto service = configBuilder.initServices(1)[0];
+    service.setName("main");
+    auto worker = service.initWorker();
+    worker.setCompatibilityDate("2023-12-18");
+    auto flags = worker.initCompatibilityFlags(2);
+    flags.set(0, compatFlagStr);
+    flags.set(1, "python_workers");
+    auto mod = worker.initModules(1)[0];
+    mod.setName("main.py");
+    mod.setPythonModule("def test():\n pass");
+    config = configBuilder.asReader();
+    configOwner = kj::mv(builder);
+    util::Autogate::initAutogate(getConfig().getAutogates());
+  }
+
   void watch() {
 #if _WIN32
     auto& w = watcher.emplace(io.win32EventPort);