From 2cd5c3639afea3f83641dea953d9495baeba5113 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                   | 11 +++++++++++
 .../server/tests/python/python-rpc/worker.py      | 15 +++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/src/pyodide/internal/workers.py b/src/pyodide/internal/workers.py
index f828fcb227e..c80b7f119b7 100644
--- a/src/pyodide/internal/workers.py
+++ b/src/pyodide/internal/workers.py
@@ -917,6 +917,17 @@ def _wrap_subclass(cls):
             continue
         setattr(cls, k, _wrap_attr(v))
 
+    # 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:
     """
diff --git a/src/workerd/server/tests/python/python-rpc/worker.py b/src/workerd/server/tests/python/python-rpc/worker.py
index 20f25ae26ee..1bd22f50c63 100644
--- a/src/workerd/server/tests/python/python-rpc/worker.py
+++ b/src/workerd/server/tests/python/python-rpc/worker.py
@@ -18,6 +18,10 @@
 
 
 class PythonRpcTester(WorkerEntrypoint):
+    def __init__(self, ctx, env):
+        assert not isinstance(env, JsProxy)
+        self.env = env
+
     async def no_args(self):
         return "hello from python"
 
@@ -37,6 +41,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):
@@ -236,3 +248,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()