From 9645f066a69ccc50783a946a86e94dd39cb0ac49 Mon Sep 17 00:00:00 2001
From: Dominik Picheta <dominik@cloudflare.com>
Date: Thu, 8 May 2025 09:58:57 +0000
Subject: [PATCH] Wrap `env` in Python DO/WorkerEntrypoint constructor.

---
 src/pyodide/internal/workers.py               | 22 +++++++++++++++++++
 .../server/tests/python/python-rpc/worker.py  | 14 ++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/src/pyodide/internal/workers.py b/src/pyodide/internal/workers.py
index 27f70efe0ba..641d6811d6c 100644
--- a/src/pyodide/internal/workers.py
+++ b/src/pyodide/internal/workers.py
@@ -892,6 +892,19 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
+def _wrap_subclass(cls):
+    # Override the class __init__ so that we can wrap the `env` in the constructor.
+    original_init = cls.__init__
+
+    def wrapped_init(self, *args, **kwargs):
+        if len(args) > 1:
+            args = (args[0], _EnvWrapper(args[1])) + args[2:]
+
+        original_init(self, *args, **kwargs)
+
+    cls.__init__ = wrapped_init
+
+
 class DurableObject:
     """
     Base class used to define a Durable Object.
@@ -901,6 +914,9 @@ def __init__(self, ctx, env):
         self.ctx = ctx
         self.env = env
 
+    def __init_subclass__(cls, **_kwargs):
+        _wrap_subclass(cls)
+
 
 class WorkerEntrypoint:
     """
@@ -911,6 +927,9 @@ def __init__(self, ctx, env):
         self.ctx = ctx
         self.env = env
 
+    def __init_subclass__(cls, **_kwargs):
+        _wrap_subclass(cls)
+
 
 class WorkflowEntrypoint:
     """
@@ -920,3 +939,6 @@ class WorkflowEntrypoint:
     def __init__(self, ctx, env):
         self.ctx = ctx
         self.env = env
+
+    def __init_subclass__(cls, **_kwargs):
+        _wrap_subclass(cls)
diff --git a/src/workerd/server/tests/python/python-rpc/worker.py b/src/workerd/server/tests/python/python-rpc/worker.py
index c33c9d39ea8..edf4cccc5ca 100644
--- a/src/workerd/server/tests/python/python-rpc/worker.py
+++ b/src/workerd/server/tests/python/python-rpc/worker.py
@@ -24,6 +24,9 @@ def __init__(self, ctx, env):
         assert self.env is not None
         assert self.ctx is not None
 
+        assert not isinstance(env, JsProxy)
+        self.env = env
+
     async def no_args(self):
         return "hello from python"
 
@@ -43,6 +46,14 @@ async def handle_request(self, req):
         assert isinstance(req, Request)
         return req
 
+    async def check_env(self):
+        # Verify that the `env` supplied to the entrypoint class is wrapped.
+        curr_date = datetime.now()
+        py_date = await self.env.PythonRpc.identity(curr_date)
+        assert isinstance(py_date, datetime)
+        assert py_date == curr_date
+        return True
+
 
 class CustomType:
     def __init__(self, x):
@@ -242,3 +253,6 @@ def my_func():
         await env.PythonRpc.identity(my_func)
     with assertRaises(TypeError):
         await env.PythonRpc.identity({"test": (1, 2, 3)})
+
+    # Verify that the `env` in the DO is correctly wrapped.
+    assert await env.PythonRpc.check_env()