Wrap env in Python DO/WorkerEntrypoint constructor.

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)

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()