Python asyncio.StreamWriter() Examples

def test_socks5_http_create_connection(event_loop):
    reader = asyncio.StreamReader(loop=event_loop)
    protocol = asyncio.StreamReaderProtocol(reader, loop=event_loop)

    transport, _ = await create_connection(
        proxy_url=SOCKS5_IPV4_URL,
        protocol_factory=lambda: protocol,
        host=HTTP_TEST_HOST,
        port=HTTP_TEST_PORT,
    )

    writer = asyncio.StreamWriter(transport, protocol, reader, event_loop)

    request = ("GET /ip HTTP/1.1\r\n"
               "Host: %s\r\n"
               "Connection: close\r\n\r\n" % HTTP_TEST_HOST)

    writer.write(request.encode())
    response = await reader.read(-1)
    assert b'200 OK' in response 

def handshake(
    remote: kademlia.Node, privkey: datatypes.PrivateKey, token: CancelToken
) -> Tuple[
    bytes, bytes, BasePreImage, BasePreImage, asyncio.StreamReader, asyncio.StreamWriter
]:  # noqa: E501
    """
    Perform the auth handshake with given remote.

    Returns the established secrets and the StreamReader/StreamWriter pair already connected to
    the remote.
    """
    use_eip8 = False
    initiator = HandshakeInitiator(remote, privkey, use_eip8, token)
    reader, writer = await initiator.connect()
    opened_connections[remote.__repr__()] = (reader, writer)
    aes_secret, mac_secret, egress_mac, ingress_mac = await _handshake(
        initiator, reader, writer, token
    )
    return aes_secret, mac_secret, egress_mac, ingress_mac, reader, writer 

def __call__(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
        """Main entry point for an interpreter session with a single client."""
        self.reader = reader
        self.writer = writer
        self.running = True

        if self.banner:
            writer.write(self.banner)
            await writer.drain()

        while self.running:
            try:
                await self.handle_one_command()
            except ConnectionResetError:
                writer.close()
                self.running = False
                break
            except Exception:
                log.exception("Exception in manhole REPL")
                self.writer.write(traceback.format_exc())
                await self.writer.drain() 

def get_directly_connected_streams(alice_extra_info: Dict[str, Any] = None,
                                   bob_extra_info: Dict[str, Any] = None,
                                   loop: asyncio.AbstractEventLoop = None) -> TConnectedStreams:
    if loop is None:
        loop = asyncio.get_event_loop()

    alice_reader = asyncio.StreamReader()
    bob_reader = asyncio.StreamReader()

    alice_transport = MemoryWriteTransport(bob_reader, extra=alice_extra_info)
    bob_transport = MemoryWriteTransport(alice_reader, extra=bob_extra_info)

    alice_protocol = MemoryProtocol()
    bob_protocol = MemoryProtocol()

    # Link the alice's writer to the bob's reader, and the bob's writer to the
    # alice's reader.
    bob_writer = asyncio.StreamWriter(bob_transport, bob_protocol, alice_reader, loop=loop)
    alice_writer = asyncio.StreamWriter(alice_transport, alice_protocol, bob_reader, loop=loop)
    return (
        (alice_reader, alice_writer),
        (bob_reader, bob_writer),
    ) 

def create_pipe_streams_pair():
    path = tempfile.mktemp()
    loop = asyncio.get_event_loop()
    server_side = asyncio.Future()

    def factory():

        def client_connected_cb(reader, writer):
            server_side.set_result((reader, writer))
        reader = asyncio.StreamReader(loop=loop)
        return asyncio.StreamReaderProtocol(reader, client_connected_cb, loop=loop)

    server = yield from loop.create_unix_server(factory, path)

    r1 = asyncio.StreamReader(loop=loop)
    protocol = asyncio.StreamReaderProtocol(r1, loop=loop)
    transport, _ = yield from loop.create_unix_connection(
        lambda: protocol, path)
    w1 = asyncio.StreamWriter(transport, protocol, r1, loop)

    r2, w2 = yield from server_side
    server.close()
    return (r1, w1), (r2, w2) 

def test_init_writer_initializes_a_nonblocking_pipe_streamwriter(
        self
    ):

        handler = AsyncStreamHandler(
            stream=self.write_pipe, level=10, formatter=Mock()
        )

        self.assertFalse(handler.initialized)

        await handler._init_writer()

        self.assertIsInstance(handler.writer, asyncio.StreamWriter)
        self.assertIsInstance(handler.writer._protocol, AiologgerProtocol)
        self.assertEqual(handler.writer.transport._pipe, self.write_pipe)
        self.assertTrue(handler.initialized)

        await handler.close() 

def _init_writer(self) -> StreamWriter:
        async with self._initialization_lock:
            if self.writer is not None:
                return self.writer

            transport, protocol = await self.loop.connect_write_pipe(
                self.protocol_class, self.stream
            )

            self.writer = StreamWriter(  # type: ignore # https://github.com/python/typeshed/pull/2719
                transport=transport,
                protocol=protocol,
                reader=None,
                loop=self.loop,
            )
            return self.writer 

def __init__(
        self,
        stream=None,
        level: Union[str, int, LogLevel] = LogLevel.NOTSET,
        formatter: Formatter = None,
        filter: Filter = None,
        *,
        loop: Optional[AbstractEventLoop] = None,
    ) -> None:
        super().__init__(loop=loop)
        if stream is None:
            stream = sys.stderr
        self.stream = stream
        self.level = level
        if formatter is None:
            formatter = Formatter()
        self.formatter: Formatter = formatter
        if filter:
            self.add_filter(filter)
        self.protocol_class = AiologgerProtocol
        self._initialization_lock = asyncio.Lock(loop=self.loop)
        self.writer: Optional[StreamWriter] = None 

def test_emit_writes_records_into_the_stream(self):
        msg = self.record.msg
        formatter = Mock(format=Mock(return_value=msg))
        writer = Mock(write=Mock(), drain=CoroutineMock())

        with patch(
            "aiologger.handlers.streams.StreamWriter", return_value=writer
        ):
            handler = AsyncStreamHandler(
                level=10, stream=self.write_pipe, formatter=formatter
            )

            await handler.emit(self.record)

            writer.write.assert_called_once_with(
                (msg + handler.terminator).encode()
            )
            writer.drain.assert_awaited_once()
            await handler.close() 

def test_emit_calls_handle_error_if_an_error_occurs(self):
        writer = Mock(write=CoroutineMock(), drain=CoroutineMock())
        with patch(
            "aiologger.handlers.streams.StreamWriter", return_value=writer
        ):
            exc = Exception("XABLAU")
            handler = AsyncStreamHandler(
                level=10,
                stream=self.write_pipe,
                formatter=Mock(format=Mock(side_effect=exc)),
            )
            with asynctest.patch.object(
                handler, "handle_error"
            ) as handle_error:
                await handler.emit(self.record)

                handle_error.assert_awaited_once_with(self.record, exc)
                writer.write.assert_not_awaited()
                writer.drain.assert_not_awaited() 

def receive_handshake(
            self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:

        try:
            try:
                await self._receive_handshake(reader, writer)
            except BaseException:
                if not reader.at_eof():
                    reader.feed_eof()
                writer.close()
                raise
        except COMMON_RECEIVE_HANDSHAKE_EXCEPTIONS as e:
            peername = writer.get_extra_info("peername")
            self.logger.debug("Could not complete handshake with %s: %s", peername, e)
        except asyncio.CancelledError:
            # This exception should just bubble.
            raise
        except Exception:
            peername = writer.get_extra_info("peername")
            self.logger.exception("Unexpected error handling handshake with %s", peername) 

def aclosing_multiple_writers(*writers: asyncio.StreamWriter):
    """Closes StreamWriters on clean exits, aborts them on exceptions.

    The "as" clause returns a set, and more StreamWriters can be added to the
    set.
    """
    writers = set(writers)
    try:
        yield writers
    except:
        for w in writers:
            w.transport.abort()
        raise
    else:
        for w in writers:
            w.close()
    finally:
        close_tasks, _ = await asyncio.wait([w.wait_closed() for w in writers])
        for t in close_tasks:
            if t.exception():
                logger.debug(
                    'wait_closed() raised exception: %r', t.exception()) 

def handle_ext_server_connection(
        upstream_host: str,
        upstream_port: int,
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
        info: adapters.ExtOrPortClientConnection,
) -> None:
    handler_logger.info('Connection received from %r', info)
    async with contexts.log_unhandled_exc(handler_logger), \
               contexts.aclosing_multiple_writers(writer) as writers:
        try:
            ureader, uwriter = await asyncio.open_connection(
                upstream_host, upstream_port)
        except OSError as e:
            handler_logger.warning(
                'Error while connecting to upstream: %r', e)
            return
        writers.add(writer)
        try:
            await relays.relay(reader, writer, ureader, uwriter)
        except OSError as e:
            handler_logger.warning('Connection from %r caught %r', info, e) 

def __init__(self, reader: StreamReader, writer: StreamWriter) -> None:
        self.writer = writer
        self.reader = reader
        self._drain_lock = asyncio.Lock() 

def __init__(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
        self.r = reader
        self.w = writer 

def _handshake(initiator: 'HandshakeInitiator', reader: asyncio.StreamReader,
                     writer: asyncio.StreamWriter,
                     ) -> Tuple[bytes, bytes, sha3.keccak_256, sha3.keccak_256]:
    """See the handshake() function above.

    This code was factored out into this helper so that we can create Peers with directly
    connected readers/writers for our tests.
    """
    initiator_nonce = keccak(os.urandom(HASH_LEN))
    auth_msg = initiator.create_auth_message(initiator_nonce)
    auth_init = initiator.encrypt_auth_message(auth_msg)

    if writer.transport.is_closing():
        raise HandshakeFailure("Error during handshake with {initiator.remote!r}. Writer closed.")

    writer.write(auth_init)

    auth_ack = await asyncio.wait_for(reader.read(ENCRYPTED_AUTH_ACK_LEN), timeout=REPLY_TIMEOUT)

    if reader.at_eof():
        # This is what happens when Parity nodes have blacklisted us
        # (https://github.com/ethereum/py-evm/issues/901).
        raise HandshakeFailure(f"{initiator.remote!r} disconnected before sending auth ack")

    ephemeral_pubkey, responder_nonce = initiator.decode_auth_ack_message(auth_ack)
    aes_secret, mac_secret, egress_mac, ingress_mac = initiator.derive_secrets(
        initiator_nonce,
        responder_nonce,
        ephemeral_pubkey,
        auth_init,
        auth_ack
    )

    return aes_secret, mac_secret, egress_mac, ingress_mac 

def receive_dial_in(reader: asyncio.StreamReader,
                          writer: asyncio.StreamWriter,
                          private_key: keys.PrivateKey,
                          p2p_handshake_params: DevP2PHandshakeParams,
                          protocol_handshakers: Sequence[HandshakerAPI[ProtocolAPI]],
                          ) -> Connection:
    transport = await Transport.receive_connection(
        reader=reader,
        writer=writer,
        private_key=private_key,
    )
    try:
        multiplexer, devp2p_receipt, protocol_receipts = await negotiate_protocol_handshakes(
            transport=transport,
            p2p_handshake_params=p2p_handshake_params,
            protocol_handshakers=protocol_handshakers,
        )
    except BaseException:
        # Note: This is one of two places where we manually handle closing the
        # reader/writer connection pair in the event of an error during the
        # peer connection and handshake process.
        # See `p2p.auth.handshake` for the other.
        try:
            await transport.close()
        except ConnectionResetError:
            transport.logger.debug("Could not wait for transport to close")
        raise

    connection = Connection(
        multiplexer=multiplexer,
        devp2p_receipt=devp2p_receipt,
        protocol_receipts=protocol_receipts,
        is_dial_out=False,
    )
    return connection 

def connect(self) -> Tuple[asyncio.StreamReader, asyncio.StreamWriter]:
        return await asyncio.wait_for(
            asyncio.open_connection(host=self.remote.address.ip, port=self.remote.address.tcp_port),
            timeout=REPLY_TIMEOUT) 

def __init__(self,
                 remote: NodeAPI,
                 private_key: datatypes.PrivateKey,
                 reader: asyncio.StreamReader,
                 writer: asyncio.StreamWriter) -> None:
        self.logger = get_logger('p2p.tools.memory_transport.MemoryTransport')
        self.remote = remote
        self.session = Session(remote)
        self._private_key = private_key
        self._reader = reader
        self._writer = writer 

def _write_data(stream: StreamWriter, data: str):
    stream.write(data.encode("UTF-8"))
    stream.close() 

def __init__(
        self,
        writer: asyncio.StreamWriter,
        connection_number: int,
        output_queue: asyncio.Queue,
        name=None,
        loop=None,
    ):
        self._logger = logging.get_named_logger(MessageWriter, name, connection_number)
        self.writer = writer
        self._queue = output_queue 

def __init__(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
        self.r = reader
        self.w = writer 

def __init__(self, host: str, port: int, channel: Channel, password: str):
        self.host = host
        self.port = port
        self.channel = channel
        self.password = password
        self.reader = None  # type: Optional[asyncio.StreamReader]
        self.writer = None  # type: Optional[asyncio.StreamWriter]
        self.logger = getLogger('connection') 

def connect(self):
        loop = self._loop
        self._closed_fut = create_future(loop=loop)
        if self._security_protocol in ["PLAINTEXT", "SASL_PLAINTEXT"]:
            ssl = None
        else:
            assert self._security_protocol in ["SSL", "SASL_SSL"]
            assert self._ssl_context is not None
            ssl = self._ssl_context
        # Create streams same as `open_connection`, but using custom protocol
        reader = asyncio.StreamReader(limit=READER_LIMIT, loop=loop)
        protocol = AIOKafkaProtocol(self._closed_fut, reader, loop=loop)
        transport, _ = await asyncio.wait_for(
            loop.create_connection(
                lambda: protocol, self.host, self.port, ssl=ssl),
            loop=loop, timeout=self._request_timeout)
        writer = asyncio.StreamWriter(transport, protocol, reader, loop)
        self._reader, self._writer, self._protocol = reader, writer, protocol

        # Start reader task.
        self._read_task = self._create_reader_task()

        # Start idle checker
        if self._max_idle_ms is not None:
            self._idle_handle = self._loop.call_soon(
                self._idle_check, weakref.ref(self))

        if self._version_hint and self._version_hint >= (0, 10):
            await self._do_version_lookup()

        if self._security_protocol in ["SASL_SSL", "SASL_PLAINTEXT"]:
            await self._do_sasl_handshake()

        return reader, writer 

def __init__(self, configuration: ConnectionConfig, **kwargs):
        """
        Initialize a TCP server connection.

        :param configuration: the configuration object.
        """
        address = cast(str, configuration.config.get("address"))
        port = cast(int, configuration.config.get("port"))
        assert address is not None and port is not None, "address and port must be set!"
        super().__init__(address, port, configuration=configuration, **kwargs)
        self._server = None  # type: Optional[AbstractServer]
        self.connections = {}  # type: Dict[str, Tuple[StreamReader, StreamWriter]]

        self._read_tasks_to_address = dict()  # type: Dict[Future, Address] 

def select_writer_from_envelope(self, envelope: Envelope) -> Optional[StreamWriter]:
        """
        Select the destination, given the envelope.

        :param envelope: the envelope to be sent.
        :return: the stream writer to communicate with the recipient. None if it cannot be determined.
        """ 

def select_writer_from_envelope(self, envelope: Envelope) -> Optional[StreamWriter]:
        """Select the destination, given the envelope."""
        return self._writer 

def __init__(self, configuration: ConnectionConfig, **kwargs):
        """
        Initialize a TCP client connection.

        :param configuration: the configuration object.
        """
        address = cast(str, configuration.config.get("address"))
        port = cast(int, configuration.config.get("port"))
        assert address is not None and port is not None, "address and port must be set!"
        super().__init__(address, port, configuration=configuration, **kwargs)
        self._reader, self._writer = (
            None,
            None,
        )  # type: Optional[StreamReader], Optional[StreamWriter] 

def sendfile(self, writer: asyncio.StreamWriter) -> int:
        """
        Read and send the file to the writer and return the number of bytes sent
        """

        if not self.is_readable():
            raise OSError('blob files cannot be read')
        with self.reader_context() as handle:
            try:
                return await self.loop.sendfile(writer.transport, handle, count=self.get_length())
            except (ConnectionError, BrokenPipeError, RuntimeError, OSError, AttributeError):
                return -1 

def __init__(self, blob_manager: 'BlobManager', response_chunk_size: int = 10000):
        self.loop = asyncio.get_event_loop()
        self.blob_manager = blob_manager
        self.server_task: asyncio.Task = None
        self.started_listening = asyncio.Event(loop=self.loop)
        self.buf = b''
        self.transport: asyncio.StreamWriter = None
        self.writer: typing.Optional['HashBlobWriter'] = None
        self.client_version: typing.Optional[int] = None
        self.descriptor: typing.Optional['StreamDescriptor'] = None
        self.sd_blob: typing.Optional['BlobFile'] = None
        self.received = []
        self.incoming = asyncio.Event(loop=self.loop)
        self.chunk_size = response_chunk_size