Add HTMLRewriter class in Python SDK

src/pyodide/internal/workers-api/src/workers/__init__.py

@@ -11,6 +11,7 @@
     FormData,
     FormDataValue,
     Headers,
+    HTMLRewriter,
     JSBody,
     Request,
     RequestInitCfProperties,
@@ -37,6 +38,7 @@
     "File",
     "FormData",
     "FormDataValue",
+    "HTMLRewriter",
     "Headers",
     "JSBody",
     "Request",

src/pyodide/internal/workers-api/src/workers/_workers.py

@@ -1366,6 +1366,105 @@ async def wrapped_run(self, event=None, step=None, /, *args, **kwargs):
     cls.run = wrapped_run
 
 
+_ELEMENT_HANDLER_METHODS = ("element", "comments", "text")
+_DOCUMENT_HANDLER_METHODS = ("doctype", "comments", "text", "end")
+
+
+def _make_js_handler(handler, method_names, proxies):
+    js_obj = Object.new()
+    for name in method_names:
+        method = getattr(handler, name, None)
+        if method is not None:
+            proxy = _pyodide_entrypoint_helper.createHandlerGuard(create_proxy(method))
+            proxies.append(proxy)
+            setattr(js_obj, name, proxy)
+    return js_obj
+
+
+class HTMLRewriter:
+    def __init__(self):
+        self._handlers: list[tuple] = []
+        # Testing only, stores the proxies created for handlers
+        self._last_handler_proxies: list[JsProxy] | None = None
+
+    def on(self, selector: str, handlers: object) -> "HTMLRewriter":
+        self._handlers.append(("element", selector, handlers))
+        return self
+
+    def onDocument(self, handlers: object) -> "HTMLRewriter":
+        self._handlers.append(("document", handlers))
+        return self
+
+    def transform(self, response: Response) -> "Response":
+        js_rewriter = js.HTMLRewriter.new()
+        handler_proxies: list[JsProxy] = []
+
+        for handler_info in self._handlers:
+            if handler_info[0] == "element":
+                _, selector, handler = handler_info
+                js_handler = _make_js_handler(
+                    handler, _ELEMENT_HANDLER_METHODS, handler_proxies
+                )
+                js_rewriter.on(selector, js_handler)
+            else:
+                _, handler = handler_info
+                js_handler = _make_js_handler(
+                    handler, _DOCUMENT_HANDLER_METHODS, handler_proxies
+                )
+                js_rewriter.onDocument(js_handler)
+
+        self._last_handler_proxies = handler_proxies
+        transformed = js_rewriter.transform(response.js_object)
+
+        if transformed.body is None:
+            return Response(transformed)
+
+        reader = transformed.body.getReader()
+
+        def cleanup():
+            for proxy in handler_proxies:
+                proxy.destroy()
+
+        async def start(controller):
+            try:
+                while True:
+                    result = await reader.read()
+                    if result.done:
+                        controller.close()
+                        break
+                    controller.enqueue(result.value)
+            except Exception as e:
+                controller.error(e)
+            finally:
+                cleanup()
+
+        async def cancel(reason):
+            try:
+                if reader:
+                    await reader.cancel(reason)
+            finally:
+                cleanup()
+
+        start_proxy = create_proxy(start)
+        cancel_proxy = create_proxy(cancel)
+        handler_proxies.append(start_proxy)
+        handler_proxies.append(cancel_proxy)
+
+        wrapped_body = js.ReadableStream.new(
+            start=start_proxy,
+            cancel=cancel_proxy,
+        )
+
+        return Response(
+            js.Response.new(
+                wrapped_body,
+                status=transformed.status,
+                statusText=transformed.statusText,
+                headers=transformed.headers,
+            )
+        )
+
+
 class DurableObject:
     """
     Base class used to define a Durable Object.

src/pyodide/python-entrypoint-helper.ts

@@ -69,13 +69,90 @@ function patchWaitUntil(ctx: {
   waitUntilPatched.add(ctx);
 }
 
+/**
+ * Creates a guard wrapper around a Python handler proxy.
+ *
+ * The guard forwards all calls to the underlying Python proxy while it's active.
+ * Once destroy() is called, the guard becomes inert - any subsequent calls
+ * return undefined instead of throwing "Object has already been destroyed" errors.
+ *
+ * This is added to prevent Python handlers being used after Python has destroyed the proxy.
+ * TODO(later): Ideally, we should control the lifetime of the proxy and destroy it when we are certain that
+ * it is no longer needed.
+ *
+ */
+export function createHandlerGuard(pythonProxy: any): any {
+  let active = true;
+
+  return new Proxy(function () {}, {
+    get(_target, prop): any {
+      if (prop === 'destroy') {
+        return () => {
+          if (active) {
+            active = false;
+            try {
+              pythonProxy.destroy();
+            } catch (_e) {
+              // Ignore errors during destroy
+            }
+          }
+        };
+      }
+
+      if (prop === '_active') {
+        return active;
+      }
+
+      // After destruction, return no-op for any method call
+      if (!active) {
+        return () => undefined;
+      }
+
+      // Forward property access to the Python proxy
+      const value = pythonProxy[prop];
+
+      // If it's a function, wrap it to handle potential async calls
+      if (typeof value === 'function') {
+        return (...args: any[]) => {
+          return value.apply(pythonProxy, args);
+        };
+      }
+
+      return value;
+    },
+
+    apply(_target, thisArg, args): any {
+      if (!active) {
+        return undefined;
+      }
+      return pythonProxy.apply(thisArg, args);
+    },
+
+    has(_target, prop): boolean {
+      if (prop === 'destroy') {
+        return true;
+      }
+
+      if (prop === '_active') {
+        return true;
+      }
+
+      if (!active) {
+        return false;
+      }
+      return prop in pythonProxy;
+    },
+  });
+}
+
 export type PyodideEntrypointHelper = {
   doAnImport: (mod: string) => Promise<any>;
   cloudflareWorkersModule: { env: any };
   cloudflareSocketsModule: any;
   workerEntrypoint: any;
   patchWaitUntil: typeof patchWaitUntil;
   patch_env_helper: (patch: unknown) => Generator<void>;
+  createHandlerGuard: typeof createHandlerGuard;
 };
 
 // Function to import JavaScript modules from Python
@@ -101,6 +178,7 @@ export async function setDoAnImport(
     workerEntrypoint,
     patchWaitUntil,
     patch_env_helper,
+    createHandlerGuard,
   };
 }
 

src/workerd/server/tests/python/pytest/pytest.wd-test

@@ -8,6 +8,7 @@ const unitTests :Workerd.Config = (
           (name = "main.py", pythonModule = embed "pytest/main.py"),
           (name = "tests/test_env.py", pythonModule = embed "pytest/tests/test_env.py"),
           (name = "tests/test_fs.py", pythonModule = embed "pytest/tests/test_fs.py"),
+          (name = "tests/test_htmlrewriter.py", pythonModule = embed "pytest/tests/test_htmlrewriter.py"),
           (name = "tests/test_import_from_javascript.py", pythonModule = embed "pytest/tests/test_import_from_javascript.py"),
           (name = "tests/test_dynlib_loading.py", pythonModule = embed "pytest/tests/test_dynlib_loading.py"),
           %PYTHON_VENDORED_MODULES%
@@ -16,6 +17,7 @@ const unitTests :Workerd.Config = (
           %PYTHON_FEATURE_FLAGS,
           "python_no_global_handlers",
           "unwrap_custom_thenables",
+          "streams_enable_constructors",
         ],
         bindings = [
           (