203 lines
8.0 KiB
Python
203 lines
8.0 KiB
Python
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import errno
|
|||
|
|
import os
|
|||
|
|
import sys
|
|||
|
|
from typing import TYPE_CHECKING
|
|||
|
|
|
|||
|
|
import trio
|
|||
|
|
|
|||
|
|
from ._abc import Stream
|
|||
|
|
from ._util import ConflictDetector, final
|
|||
|
|
|
|||
|
|
if TYPE_CHECKING:
|
|||
|
|
from typing import Final as FinalType
|
|||
|
|
|
|||
|
|
assert not TYPE_CHECKING or sys.platform != "win32"
|
|||
|
|
|
|||
|
|
if os.name != "posix":
|
|||
|
|
# We raise an error here rather than gating the import in lowlevel.py
|
|||
|
|
# in order to keep jedi static analysis happy.
|
|||
|
|
raise ImportError
|
|||
|
|
|
|||
|
|
# XX TODO: is this a good number? who knows... it does match the default Linux
|
|||
|
|
# pipe capacity though.
|
|||
|
|
DEFAULT_RECEIVE_SIZE: FinalType = 65536
|
|||
|
|
|
|||
|
|
|
|||
|
|
class _FdHolder:
|
|||
|
|
# This class holds onto a raw file descriptor, in non-blocking mode, and
|
|||
|
|
# is responsible for managing its lifecycle. In particular, it's
|
|||
|
|
# responsible for making sure it gets closed, and also for tracking
|
|||
|
|
# whether it's been closed.
|
|||
|
|
#
|
|||
|
|
# The way we track closure is to set the .fd field to -1, discarding the
|
|||
|
|
# original value. You might think that this is a strange idea, since it
|
|||
|
|
# overloads the same field to do two different things. Wouldn't it be more
|
|||
|
|
# natural to have a dedicated .closed field? But that would be more
|
|||
|
|
# error-prone. Fds are represented by small integers, and once an fd is
|
|||
|
|
# closed, its integer value may be reused immediately. If we accidentally
|
|||
|
|
# used the old fd after being closed, we might end up doing something to
|
|||
|
|
# another unrelated fd that happened to get assigned the same integer
|
|||
|
|
# value. By throwing away the integer value immediately, it becomes
|
|||
|
|
# impossible to make this mistake – we'll just get an EBADF.
|
|||
|
|
#
|
|||
|
|
# (This trick was copied from the stdlib socket module.)
|
|||
|
|
fd: int
|
|||
|
|
|
|||
|
|
def __init__(self, fd: int) -> None:
|
|||
|
|
# make sure self.fd is always initialized to *something*, because even
|
|||
|
|
# if we error out here then __del__ will run and access it.
|
|||
|
|
self.fd = -1
|
|||
|
|
if not isinstance(fd, int):
|
|||
|
|
raise TypeError("file descriptor must be an int")
|
|||
|
|
self.fd = fd
|
|||
|
|
# Store original state, and ensure non-blocking mode is enabled
|
|||
|
|
self._original_is_blocking = os.get_blocking(fd)
|
|||
|
|
os.set_blocking(fd, False)
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def closed(self) -> bool:
|
|||
|
|
return self.fd == -1
|
|||
|
|
|
|||
|
|
def _raw_close(self) -> None:
|
|||
|
|
# This doesn't assume it's in a Trio context, so it can be called from
|
|||
|
|
# __del__. You should never call it from Trio context, because it
|
|||
|
|
# skips calling notify_fd_close. But from __del__, skipping that is
|
|||
|
|
# OK, because notify_fd_close just wakes up other tasks that are
|
|||
|
|
# waiting on this fd, and those tasks hold a reference to this object.
|
|||
|
|
# So if __del__ is being called, we know there aren't any tasks that
|
|||
|
|
# need to be woken.
|
|||
|
|
if self.closed:
|
|||
|
|
return
|
|||
|
|
fd = self.fd
|
|||
|
|
self.fd = -1
|
|||
|
|
os.set_blocking(fd, self._original_is_blocking)
|
|||
|
|
os.close(fd)
|
|||
|
|
|
|||
|
|
def __del__(self) -> None:
|
|||
|
|
self._raw_close()
|
|||
|
|
|
|||
|
|
def close(self) -> None:
|
|||
|
|
if not self.closed:
|
|||
|
|
trio.lowlevel.notify_closing(self.fd)
|
|||
|
|
self._raw_close()
|
|||
|
|
|
|||
|
|
|
|||
|
|
@final
|
|||
|
|
class FdStream(Stream):
|
|||
|
|
"""
|
|||
|
|
Represents a stream given the file descriptor to a pipe, TTY, etc.
|
|||
|
|
|
|||
|
|
*fd* must refer to a file that is open for reading and/or writing and
|
|||
|
|
supports non-blocking I/O (pipes and TTYs will work, on-disk files probably
|
|||
|
|
not). The returned stream takes ownership of the fd, so closing the stream
|
|||
|
|
will close the fd too. As with `os.fdopen`, you should not directly use
|
|||
|
|
an fd after you have wrapped it in a stream using this function.
|
|||
|
|
|
|||
|
|
To be used as a Trio stream, an open file must be placed in non-blocking
|
|||
|
|
mode. Unfortunately, this impacts all I/O that goes through the
|
|||
|
|
underlying open file, including I/O that uses a different
|
|||
|
|
file descriptor than the one that was passed to Trio. If other threads
|
|||
|
|
or processes are using file descriptors that are related through `os.dup`
|
|||
|
|
or inheritance across `os.fork` to the one that Trio is using, they are
|
|||
|
|
unlikely to be prepared to have non-blocking I/O semantics suddenly
|
|||
|
|
thrust upon them. For example, you can use
|
|||
|
|
``FdStream(os.dup(sys.stdin.fileno()))`` to obtain a stream for reading
|
|||
|
|
from standard input, but it is only safe to do so with heavy caveats: your
|
|||
|
|
stdin must not be shared by any other processes, and you must not make any
|
|||
|
|
calls to synchronous methods of `sys.stdin` until the stream returned by
|
|||
|
|
`FdStream` is closed. See `issue #174
|
|||
|
|
<https://github.com/python-trio/trio/issues/174>`__ for a discussion of the
|
|||
|
|
challenges involved in relaxing this restriction.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
fd (int): The fd to be wrapped.
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
A new `FdStream` object.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self, fd: int) -> None:
|
|||
|
|
self._fd_holder = _FdHolder(fd)
|
|||
|
|
self._send_conflict_detector = ConflictDetector(
|
|||
|
|
"another task is using this stream for send"
|
|||
|
|
)
|
|||
|
|
self._receive_conflict_detector = ConflictDetector(
|
|||
|
|
"another task is using this stream for receive"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
async def send_all(self, data: bytes) -> None:
|
|||
|
|
with self._send_conflict_detector:
|
|||
|
|
# have to check up front, because send_all(b"") on a closed pipe
|
|||
|
|
# should raise
|
|||
|
|
if self._fd_holder.closed:
|
|||
|
|
raise trio.ClosedResourceError("file was already closed")
|
|||
|
|
await trio.lowlevel.checkpoint()
|
|||
|
|
length = len(data)
|
|||
|
|
# adapted from the SocketStream code
|
|||
|
|
with memoryview(data) as view:
|
|||
|
|
sent = 0
|
|||
|
|
while sent < length:
|
|||
|
|
with view[sent:] as remaining:
|
|||
|
|
try:
|
|||
|
|
sent += os.write(self._fd_holder.fd, remaining)
|
|||
|
|
except BlockingIOError:
|
|||
|
|
await trio.lowlevel.wait_writable(self._fd_holder.fd)
|
|||
|
|
except OSError as e:
|
|||
|
|
if e.errno == errno.EBADF:
|
|||
|
|
raise trio.ClosedResourceError(
|
|||
|
|
"file was already closed"
|
|||
|
|
) from None
|
|||
|
|
else:
|
|||
|
|
raise trio.BrokenResourceError from e
|
|||
|
|
|
|||
|
|
async def wait_send_all_might_not_block(self) -> None:
|
|||
|
|
with self._send_conflict_detector:
|
|||
|
|
if self._fd_holder.closed:
|
|||
|
|
raise trio.ClosedResourceError("file was already closed")
|
|||
|
|
try:
|
|||
|
|
await trio.lowlevel.wait_writable(self._fd_holder.fd)
|
|||
|
|
except BrokenPipeError as e:
|
|||
|
|
# kqueue: raises EPIPE on wait_writable instead
|
|||
|
|
# of sending, which is annoying
|
|||
|
|
raise trio.BrokenResourceError from e
|
|||
|
|
|
|||
|
|
async def receive_some(self, max_bytes: int | None = None) -> bytes:
|
|||
|
|
with self._receive_conflict_detector:
|
|||
|
|
if max_bytes is None:
|
|||
|
|
max_bytes = DEFAULT_RECEIVE_SIZE
|
|||
|
|
else:
|
|||
|
|
if not isinstance(max_bytes, int):
|
|||
|
|
raise TypeError("max_bytes must be integer >= 1")
|
|||
|
|
if max_bytes < 1:
|
|||
|
|
raise ValueError("max_bytes must be integer >= 1")
|
|||
|
|
|
|||
|
|
await trio.lowlevel.checkpoint()
|
|||
|
|
while True:
|
|||
|
|
try:
|
|||
|
|
data = os.read(self._fd_holder.fd, max_bytes)
|
|||
|
|
except BlockingIOError:
|
|||
|
|
await trio.lowlevel.wait_readable(self._fd_holder.fd)
|
|||
|
|
except OSError as exc:
|
|||
|
|
if exc.errno == errno.EBADF:
|
|||
|
|
raise trio.ClosedResourceError(
|
|||
|
|
"file was already closed"
|
|||
|
|
) from None
|
|||
|
|
else:
|
|||
|
|
raise trio.BrokenResourceError from exc
|
|||
|
|
else:
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
return data
|
|||
|
|
|
|||
|
|
def close(self) -> None:
|
|||
|
|
self._fd_holder.close()
|
|||
|
|
|
|||
|
|
async def aclose(self) -> None:
|
|||
|
|
self.close()
|
|||
|
|
await trio.lowlevel.checkpoint()
|
|||
|
|
|
|||
|
|
def fileno(self) -> int:
|
|||
|
|
return self._fd_holder.fd
|