Lazy import either 'anyio' or 'trio'

httpcore/_synchronization.py

@@ -2,17 +2,44 @@
 from types import TracebackType
 from typing import Optional, Type
 
-import anyio
+import sniffio
 
 from ._exceptions import ExceptionMapping, PoolTimeout, map_exceptions
 
+# Our async synchronization primatives use either 'anyio' or 'trio' depending
+# on if they're running under asyncio or trio.
+#
+# We take care to only lazily import whichever of these two we need.
+
 
 class AsyncLock:
  def __init__(self) -> None:
- self._lock = anyio.Lock()
+ self._backend = ""
+
+ def setup(self) -> None:
+ """
+ Detect if we're running under 'asyncio' or 'trio' and create
+ a lock with the correct implementation.
+ """
+ self._backend = sniffio.current_async_library()
+ if self._backend == "trio":
+ import trio
+
+ self._trio_lock = trio.Lock()
+ else:
+ import anyio
+
+ self._anyio_lock = anyio.Lock()
 
  async def __aenter__(self) -> "AsyncLock":
- await self._lock.acquire()
+ if not self._backend:
+ self.setup()
+
+ if self._backend == "trio":
+ await self._trio_lock.acquire()
+ else:
+ await self._anyio_lock.acquire()
+
  return self
 
  async def __aexit__(
@@ -21,32 +48,102 @@ async def __aexit__(
  exc_value: Optional[BaseException] = None,
  traceback: Optional[TracebackType] = None,
  ) -> None:
- self._lock.release()
+ if self._backend == "trio":
+ self._trio_lock.release()
+ else:
+ self._anyio_lock.release()
 
 
 class AsyncEvent:
  def __init__(self) -> None:
- self._event = anyio.Event()
+ self._backend = ""
+
+ def setup(self) -> None:
+ """
+ Detect if we're running under 'asyncio' or 'trio' and create
+ a lock with the correct implementation.
+ """
+ self._backend = sniffio.current_async_library()
+ if self._backend == "trio":
+ import trio
+
+ self._trio_event = trio.Event()
+ else:
+ import anyio
+
+ self._anyio_event = anyio.Event()
 
  def set(self) -> None:
- self._event.set()
+ if not self._backend:
+ self.setup()
+
+ if self._backend == "trio":
+ self._trio_event.set()
+ else:
+ self._anyio_event.set()
 
  async def wait(self, timeout: Optional[float] = None) -> None:
- exc_map: ExceptionMapping = {TimeoutError: PoolTimeout}
- with map_exceptions(exc_map):
- with anyio.fail_after(timeout):
- await self._event.wait()
+ if not self._backend:
+ self.setup()
+
+ if self._backend == "trio":
+ import trio
+
+ trio_exc_map: ExceptionMapping = {trio.TooSlowError: PoolTimeout}
+ timeout_or_inf = float("inf") if timeout is None else timeout
+ with map_exceptions(trio_exc_map):
+ with trio.fail_after(timeout_or_inf):
+ await self._trio_event.wait()
+ else:
+ import anyio
+
+ anyio_exc_map: ExceptionMapping = {TimeoutError: PoolTimeout}
+ with map_exceptions(anyio_exc_map):
+ with anyio.fail_after(timeout):
+ await self._anyio_event.wait()
 
 
 class AsyncSemaphore:
  def __init__(self, bound: int) -> None:
- self._semaphore = anyio.Semaphore(initial_value=bound, max_value=bound)
+ self._bound = bound
+ self._backend = ""
+
+ def setup(self) -> None:
+ """
+ Detect if we're running under 'asyncio' or 'trio' and create
+ a semaphore with the correct implementation.
+ """
+ self._backend = sniffio.current_async_library()
+ if self._backend == "trio":
+ import trio
+
+ self._trio_semaphore = trio.Semaphore(
+ initial_value=self._bound, max_value=self._bound
+ )
+ else:
+ import anyio
+
+ self._anyio_semaphore = anyio.Semaphore(
+ initial_value=self._bound, max_value=self._bound
+ )
 
  async def acquire(self) -> None:
- await self._semaphore.acquire()
+ if not self._backend:
+ self.setup()
+
+ if self._backend == "trio":
+ await self._trio_semaphore.acquire()
+ else:
+ await self._anyio_semaphore.acquire()
 
  async def release(self) -> None:
- self._semaphore.release()
+ if self._backend == "trio":
+ self._trio_semaphore.release()
+ else:
+ self._anyio_semaphore.release()
+
+
+# Our thread-based synchronization primitives...
 
 
 class Lock: