Repository: django/asgiref
Pull request #244 (2021-03-04 18:29:20)

Parameterize executor
There are cases where it may be desirable to run certain sync_to_async callables in their own executor. For example, a user might want to allow a method to be run in small thread pool without affecting the loop's default thread pool. Or, a custom executor may be passed in to gather metrics about calls for certain decorated methods.


Patch:
diff --git a/asgiref/sync.py b/asgiref/sync.py
index a8d91ea..22b5307 100644
--- a/asgiref/sync.py
+++ b/asgiref/sync.py
@@ -5,7 +5,7 @@ import sys
 import threading
 import weakref
 from concurrent.futures import Future, ThreadPoolExecutor
-from typing import Dict
+from typing import Any, Callable, Dict, Optional, Union
 
 from .current_thread_executor import CurrentThreadExecutor
 from .local import Local
@@ -293,6 +293,10 @@ class SyncToAsync:
     outermost), this will just be the main thread. This is achieved by idling
     with a CurrentThreadExecutor while AsyncToSync is blocking its sync parent,
     rather than just blocking.
+
+    If executor is passed in, that will be used instead of the loop's default executor.
+    In order to pass in an executor, thread_sensitive must be set to False, otherwise
+    a TypeError will be raised.
     """
 
     # If they've set ASGI_THREADS, update the default asyncio executor for now
@@ -326,15 +330,23 @@ class SyncToAsync:
         weakref.WeakKeyDictionary()
     )
 
-    def __init__(self, func, thread_sensitive=True):
+    def __init__(
+        self,
+        func: Callable[..., Any],
+        thread_sensitive: bool = True,
+        executor: Optional["ThreadPoolExecutor"] = None,
+    ) -> None:
         if not callable(func) or asyncio.iscoroutinefunction(func):
             raise TypeError("sync_to_async can only be applied to sync functions.")
         self.func = func
         functools.update_wrapper(self, func)
         self._thread_sensitive = thread_sensitive
-        self._is_coroutine = asyncio.coroutines._is_coroutine
+        self._is_coroutine = asyncio.coroutines._is_coroutine  # type: ignore
+        if thread_sensitive and executor is not None:
+            raise TypeError("executor must not be set when thread_sensitive is True")
+        self._executor = executor
         try:
-            self.__self__ = func.__self__
+            self.__self__ = func.__self__  # type: ignore
         except AttributeError:
             pass
 
@@ -364,7 +376,8 @@ class SyncToAsync:
                 # Otherwise, we run it in a fixed single thread
                 executor = self.single_thread_executor
         else:
-            executor = None  # Use default
+            # Use the passed in executor, or the loop's default if it is None
+            executor = self._executor
 
         if contextvars is not None:
             context = contextvars.copy_context()
@@ -456,13 +469,19 @@ class SyncToAsync:
 async_to_sync = AsyncToSync
 
 
-def sync_to_async(func=None, thread_sensitive=True):
+def sync_to_async(
+    func: Optional[Callable[..., Any]] = None,
+    thread_sensitive: bool = True,
+    executor: Optional["ThreadPoolExecutor"] = None,
+) -> Union[SyncToAsync, Callable[[Callable[..., Any]], SyncToAsync]]:
     if func is None:
         return lambda f: SyncToAsync(
             f,
             thread_sensitive=thread_sensitive,
+            executor=executor,
         )
     return SyncToAsync(
         func,
         thread_sensitive=thread_sensitive,
+        executor=executor,
     )


Test patch:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 3bc8d80..a531af3 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -3,7 +3,7 @@ name: Tests
 on:
   push:
    branches:
-   - master
+   - main
   pull_request:
 
 jobs:
diff --git a/tests/test_sync.py b/tests/test_sync.py
index b0882a1..0aa336c 100644
--- a/tests/test_sync.py
+++ b/tests/test_sync.py
@@ -583,3 +583,40 @@ async def test_multiprocessing():
         return test_queue.get(True, 1)
 
     assert await sync_to_async(fork_first)() == 42
+
+
+@pytest.mark.asyncio
+async def test_sync_to_async_uses_executor():
+    """
+    Tests that SyncToAsync uses the passed in executor correctly.
+    """
+
+    class CustomExecutor:
+        def __init__(self):
+            self.executor = ThreadPoolExecutor(max_workers=1)
+            self.times_submit_called = 0
+
+        def submit(self, callable_, *args, **kwargs):
+            self.times_submit_called += 1
+            return self.executor.submit(callable_, *args, **kwargs)
+
+    expected_result = "expected_result"
+
+    def sync_func():
+        return expected_result
+
+    custom_executor = CustomExecutor()
+    async_function = sync_to_async(
+        sync_func, thread_sensitive=False, executor=custom_executor
+    )
+    actual_result = await async_function()
+    assert actual_result == expected_result
+    assert custom_executor.times_submit_called == 1
+
+    pytest.raises(
+        TypeError,
+        sync_to_async,
+        sync_func,
+        thread_sensitive=True,
+        executor=custom_executor,
+    )