--- /dev/null
+import functools
+import math
+import ssl
+import typing
+from types import TracebackType
+
+import trio
+
+from ..config import PoolLimits, TimeoutConfig
+from ..exceptions import ConnectTimeout, PoolTimeout, ReadTimeout, WriteTimeout
+from .base import (
+ BaseBackgroundManager,
+ BaseEvent,
+ BasePoolSemaphore,
+ BaseQueue,
+ BaseTCPStream,
+ ConcurrencyBackend,
+ TimeoutFlag,
+)
+
+
+def _or_inf(value: typing.Optional[float]) -> float:
+ return value if value is not None else float("inf")
+
+
+class TCPStream(BaseTCPStream):
+ def __init__(
+ self,
+ stream: typing.Union[trio.SocketStream, trio.SSLStream],
+ timeout: TimeoutConfig,
+ ) -> None:
+ self.stream = stream
+ self.timeout = timeout
+ self.write_buffer = b""
+ self.write_lock = trio.Lock()
+
+ def get_http_version(self) -> str:
+ if not isinstance(self.stream, trio.SSLStream):
+ return "HTTP/1.1"
+
+ ident = self.stream.selected_alpn_protocol()
+ if ident is None:
+ return "HTTP/1.1"
+
+ return "HTTP/2" if ident == "h2" else "HTTP/1.1"
+
+ async def read(
+ self, n: int, timeout: TimeoutConfig = None, flag: TimeoutFlag = None
+ ) -> bytes:
+ if timeout is None:
+ timeout = self.timeout
+
+ while True:
+ # Check our flag at the first possible moment, and use a fine
+ # grained retry loop if we're not yet in read-timeout mode.
+ should_raise = flag is None or flag.raise_on_read_timeout
+ read_timeout = _or_inf(timeout.read_timeout if should_raise else 0.01)
+
+ with trio.move_on_after(read_timeout):
+ return await self.stream.receive_some(max_bytes=n)
+
+ if should_raise:
+ raise ReadTimeout() from None
+
+ def is_connection_dropped(self) -> bool:
+ # Adapted from: https://github.com/encode/httpx/pull/143#issuecomment-515202982
+ stream = self.stream
+
+ # Peek through any SSLStream wrappers to get the underlying SocketStream.
+ while hasattr(stream, "transport_stream"):
+ stream = stream.transport_stream
+ assert isinstance(stream, trio.SocketStream)
+
+ # Counter-intuitively, what we really want to know here is whether the socket is
+ # *readable*, i.e. whether it would return immediately with empty bytes if we
+ # called `.recv()` on it, indicating that the other end has closed the socket.
+ # See: https://github.com/encode/httpx/pull/143#issuecomment-515181778
+ return stream.socket.is_readable()
+
+ def write_no_block(self, data: bytes) -> None:
+ self.write_buffer += data
+
+ async def write(
+ self, data: bytes, timeout: TimeoutConfig = None, flag: TimeoutFlag = None
+ ) -> None:
+ if self.write_buffer:
+ previous_data = self.write_buffer
+ # Reset before recursive call, otherwise we'll go through
+ # this branch indefinitely.
+ self.write_buffer = b""
+ try:
+ await self.write(previous_data, timeout=timeout, flag=flag)
+ except WriteTimeout:
+ self.writer_buffer = previous_data
+ raise
+
+ if not data:
+ return
+
+ if timeout is None:
+ timeout = self.timeout
+
+ write_timeout = _or_inf(timeout.write_timeout)
+
+ while True:
+ with trio.move_on_after(write_timeout):
+ async with self.write_lock:
+ await self.stream.send_all(data)
+ break
+ # We check our flag at the first possible moment, in order to
+ # allow us to suppress write timeouts, if we've since
+ # switched over to read-timeout mode.
+ should_raise = flag is None or flag.raise_on_write_timeout
+ if should_raise:
+ raise WriteTimeout() from None
+
+ async def close(self) -> None:
+ await self.stream.aclose()
+
+
+class PoolSemaphore(BasePoolSemaphore):
+ def __init__(self, pool_limits: PoolLimits):
+ self.pool_limits = pool_limits
+
+ @property
+ def semaphore(self) -> typing.Optional[trio.Semaphore]:
+ if not hasattr(self, "_semaphore"):
+ max_connections = self.pool_limits.hard_limit
+ if max_connections is None:
+ self._semaphore = None
+ else:
+ self._semaphore = trio.Semaphore(
+ max_connections, max_value=max_connections
+ )
+ return self._semaphore
+
+ async def acquire(self) -> None:
+ if self.semaphore is None:
+ return
+
+ timeout = _or_inf(self.pool_limits.pool_timeout)
+
+ with trio.move_on_after(timeout):
+ await self.semaphore.acquire()
+ return
+
+ raise PoolTimeout()
+
+ def release(self) -> None:
+ if self.semaphore is None:
+ return
+
+ self.semaphore.release()
+
+
+class TrioBackend(ConcurrencyBackend):
+ async def open_tcp_stream(
+ self,
+ hostname: str,
+ port: int,
+ ssl_context: typing.Optional[ssl.SSLContext],
+ timeout: TimeoutConfig,
+ ) -> TCPStream:
+ connect_timeout = _or_inf(timeout.connect_timeout)
+
+ with trio.move_on_after(connect_timeout) as cancel_scope:
+ stream: trio.SocketStream = await trio.open_tcp_stream(hostname, port)
+ if ssl_context is not None:
+ stream = trio.SSLStream(stream, ssl_context, server_hostname=hostname)
+ await stream.do_handshake()
+
+ if cancel_scope.cancelled_caught:
+ raise ConnectTimeout()
+
+ return TCPStream(stream=stream, timeout=timeout)
+
+ async def run_in_threadpool(
+ self, func: typing.Callable, *args: typing.Any, **kwargs: typing.Any
+ ) -> typing.Any:
+ return await trio.to_thread.run_sync(
+ functools.partial(func, **kwargs) if kwargs else func, *args
+ )
+
+ def run(
+ self, coroutine: typing.Callable, *args: typing.Any, **kwargs: typing.Any
+ ) -> typing.Any:
+ return trio.run(
+ functools.partial(coroutine, **kwargs) if kwargs else coroutine, *args
+ )
+
+ def get_semaphore(self, limits: PoolLimits) -> BasePoolSemaphore:
+ return PoolSemaphore(limits)
+
+ def create_queue(self, max_size: int) -> BaseQueue:
+ return Queue(max_size=max_size)
+
+ def create_event(self) -> BaseEvent:
+ return Event()
+
+ def background_manager(
+ self, coroutine: typing.Callable, *args: typing.Any
+ ) -> "BackgroundManager":
+ return BackgroundManager(coroutine, *args)
+
+
+class Queue(BaseQueue):
+ def __init__(self, max_size: int) -> None:
+ self.send_channel, self.receive_channel = trio.open_memory_channel(math.inf)
+
+ async def get(self) -> typing.Any:
+ return await self.receive_channel.receive()
+
+ async def put(self, value: typing.Any) -> None:
+ await self.send_channel.send(value)
+
+
+class Event(BaseEvent):
+ def __init__(self) -> None:
+ self._event = trio.Event()
+
+ def set(self) -> None:
+ self._event.set()
+
+ def is_set(self) -> bool:
+ return self._event.is_set()
+
+ async def wait(self) -> None:
+ await self._event.wait()
+
+ def clear(self) -> None:
+ # trio.Event.clear() was deprecated in Trio 0.12.
+ # https://github.com/python-trio/trio/issues/637
+ self._event = trio.Event()
+
+
+class BackgroundManager(BaseBackgroundManager):
+ def __init__(self, coroutine: typing.Callable, *args: typing.Any) -> None:
+ self.coroutine = coroutine
+ self.args = args
+ self.nursery_manager = trio.open_nursery()
+ self.nursery: typing.Optional[trio.Nursery] = None
+
+ async def __aenter__(self) -> "BackgroundManager":
+ self.nursery = await self.nursery_manager.__aenter__()
+ self.nursery.start_soon(self.coroutine, *self.args)
+ return self
+
+ async def __aexit__(
+ self,
+ exc_type: typing.Type[BaseException] = None,
+ exc_value: BaseException = None,
+ traceback: TracebackType = None,
+ ) -> None:
+ assert self.nursery is not None
+ await self.nursery_manager.__aexit__(exc_type, exc_value, traceback)
projects / thirdparty / httpx.git / commitdiff
