diff --git a/temporalio/worker/workflow_sandbox/_restrictions.py b/temporalio/worker/workflow_sandbox/_restrictions.py index 78b7a0363..34f9efee2 100644 --- a/temporalio/worker/workflow_sandbox/_restrictions.py +++ b/temporalio/worker/workflow_sandbox/_restrictions.py @@ -860,6 +860,8 @@ def set_on_proxy(self, v: _RestrictedProxy) -> None: class _RestrictedProxyLookup: + bind_func: Callable[[_RestrictedProxy, Any], Callable[..., Any]] | None + def __init__( self, access_func: Callable | None = None, @@ -951,12 +953,12 @@ def __init__( def bind_f(instance: _RestrictedProxy, obj: Any) -> Callable: def i_op(self: Any, other: Any) -> _RestrictedProxy: - f(self, other) # type: ignore + access_func(self, other) # type: ignore return instance return i_op.__get__(obj, type(obj)) # type: ignore - self.bind_f = bind_f + self.bind_func = bind_f _OpF = TypeVar("_OpF", bound=Callable[..., Any]) diff --git a/tests/worker/workflow_sandbox/test_restrictions.py b/tests/worker/workflow_sandbox/test_restrictions.py index bd1aaf749..ec001ba19 100644 --- a/tests/worker/workflow_sandbox/test_restrictions.py +++ b/tests/worker/workflow_sandbox/test_restrictions.py @@ -75,6 +75,21 @@ def test_restricted_proxy_dunder_methods(): assert f"{restricted_path_obj}" == expected_path +def test_restricted_proxy_in_place_operations_preserve_proxy(): + restricted_list = _RestrictedProxy( + "list", + list, + RestrictionContext(), + SandboxMatcher(), + ) + values = restricted_list([1]) + + values += [2] + + assert type(values) is _RestrictedProxy + assert RestrictionContext.unwrap_if_proxied(values) == [1, 2] + + def test_workflow_sandbox_restricted_proxy(): obj_class = _RestrictedProxy( "RestrictableObject",