added a custom context parameter for the sync_to_async

Author: Arfey

asgiref/sync.py

@@ -424,6 +424,7 @@ def __init__(
         func: Callable[_P, _R],
         thread_sensitive: bool = True,
         executor: Optional["ThreadPoolExecutor"] = None,
+        context: Optional[contextvars.Context] = None,
     ) -> None:
         if (
             not callable(func)
@@ -432,6 +433,7 @@ def __init__(
         ):
             raise TypeError("sync_to_async can only be applied to sync functions.")
         self.func = func
+        self.context = context
         functools.update_wrapper(self, func)
         self._thread_sensitive = thread_sensitive
         markcoroutinefunction(self)
@@ -480,7 +482,7 @@ async def __call__(self, *args: _P.args, **kwargs: _P.kwargs) -> _R:
             # Use the passed in executor, or the loop's default if it is None
             executor = self._executor
 
-        context = contextvars.copy_context()
+        context = contextvars.copy_context() if self.context is None else self.context
         child = functools.partial(self.func, *args, **kwargs)
         func = context.run
         task_context: List[asyncio.Task[Any]] = []
@@ -518,7 +520,8 @@ async def __call__(self, *args: _P.args, **kwargs: _P.kwargs) -> _R:
                 exec_coro.cancel()
             ret = await exec_coro
         finally:
-            _restore_context(context)
+            if self.context is None:
+                _restore_context(context)
             self.deadlock_context.set(False)
 
         return ret
@@ -611,6 +614,7 @@ def sync_to_async(
     *,
     thread_sensitive: bool = True,
     executor: Optional["ThreadPoolExecutor"] = None,
+    context: Optional[contextvars.Context] = None,
 ) -> Callable[[Callable[_P, _R]], Callable[_P, Coroutine[Any, Any, _R]]]:
     ...
 
@@ -621,6 +625,7 @@ def sync_to_async(
     *,
     thread_sensitive: bool = True,
     executor: Optional["ThreadPoolExecutor"] = None,
+    context: Optional[contextvars.Context] = None,
 ) -> Callable[_P, Coroutine[Any, Any, _R]]:
     ...
 
@@ -630,6 +635,7 @@ def sync_to_async(
     *,
     thread_sensitive: bool = True,
     executor: Optional["ThreadPoolExecutor"] = None,
+    context: Optional[contextvars.Context] = None,
 ) -> Union[
     Callable[[Callable[_P, _R]], Callable[_P, Coroutine[Any, Any, _R]]],
     Callable[_P, Coroutine[Any, Any, _R]],
@@ -639,9 +645,11 @@ def sync_to_async(
             f,
             thread_sensitive=thread_sensitive,
             executor=executor,
+            context=context,
         )
     return SyncToAsync(
         func,
         thread_sensitive=thread_sensitive,
         executor=executor,
+        context=context,
     )

tests/test_sync_contextvars.py

@@ -1,5 +1,6 @@
 import asyncio
 import contextvars
+import sys
 import threading
 import time
 
@@ -55,6 +56,75 @@ def sync_function():
     assert foo.get() == "baz"
 
 
+@pytest.mark.asyncio
+async def test_sync_to_async_contextvars_with_custom_context():
+    """
+    Test that passing a custom context to `sync_to_async` ensures that changes to
+    context variables within the synchronous function are isolated to the
+    provided context and do not affect the caller's context. Specifically,
+    verifies that modifications to a context variable inside the
+    sync function are reflected only in the custom context and not in the
+    outer context.
+    """
+    # Define sync function
+    def sync_function():
+        time.sleep(1)
+        assert foo.get() == "bar"
+        foo.set("baz")
+        return 42
+
+    # Ensure outermost detection works
+    # Wrap it
+    foo.set("bar")
+    context = contextvars.copy_context()
+    async_function = sync_to_async(sync_function, context=context)
+    assert await async_function() == 42
+
+    # verify that the current context remains unchanged
+    assert foo.get() == "bar"
+
+    # verify that the custom context reflects the changes made within the
+    # sync function
+    assert context.get(foo) == "baz"
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(sys.version_info < (3, 11), reason="requires python3.11")
+async def test_sync_to_async_contextvars_with_custom_context_and_parallel_tasks():
+    """
+    Test that using a custom context with `sync_to_async` and asyncio tasks
+    isolates contextvars changes, leaving the original context unchanged and
+    reflecting all modifications in the custom context.
+    """
+    # Ensure outermost detection works
+    # Wrap it
+    foo.set("")
+
+    def sync_function():
+        foo.set(foo.get() + "1")
+        return 1
+
+    async def async_function():
+        foo.set(foo.get() + "1")
+        return 1
+
+    context = contextvars.copy_context()
+
+    await asyncio.gather(
+        sync_to_async(sync_function, context=context)(),
+        sync_to_async(sync_function, context=context)(),
+        asyncio.create_task(async_function(), context=context),
+        asyncio.create_task(async_function(), context=context),
+    )
+
+    # verify that the current context remains unchanged
+    assert foo.get() == ""
+
+    # verify that the custom context reflects the changes made within the
+    # sync function
+    assert context.get(foo) == "1111"
+
+
 def test_async_to_sync_contextvars():
     """
     Tests to make sure that contextvars from the calling context are