def recv(self, flags=None, reconnect=True) -> bytearray:
"""
Receive data from the socket.
:param flags: (optional) OS-specific flags.
:param reconnect: (optional) reconnect on failure, default True.
"""
if self.closed:
raise SocketError('Attempt to use closed connection.')
kwargs = {}
if flags is not None:
kwargs['flags'] = flags
data = bytearray(1024)
buffer = memoryview(data)
bytes_total_received, bytes_to_receive = 0, 0
while True:
try:
bytes_received = self._socket.recv_into(
buffer, len(buffer), **kwargs)
if bytes_received == 0:
raise SocketError('Connection broken.')
bytes_total_received += bytes_received
except connection_errors:
self.failed = True
if reconnect:
self.reconnect()
raise
if bytes_total_received < 4:
continue
elif bytes_to_receive == 0:
response_len = int.from_bytes(data[0:4], PROTOCOL_BYTE_ORDER)
bytes_to_receive = response_len
if response_len + 4 > len(data):
buffer.release()
data.extend(bytearray(response_len + 4 - len(data)))
buffer = memoryview(data)[bytes_total_received:]
continue
if bytes_total_received >= bytes_to_receive:
buffer.release()
break
buffer = buffer[bytes_received:]
return data
async def _recv(self, reconnect=True) -> bytearray:
if self.closed:
raise SocketError('Attempt to use closed connection.')
data = bytearray(1024)
buffer = memoryview(data)
bytes_total_received, bytes_to_receive = 0, 0
while True:
try:
chunk = await self._reader.read(len(buffer))
bytes_received = len(chunk)
if bytes_received == 0:
raise SocketError('Connection broken.')
buffer[0:bytes_received] = chunk
bytes_total_received += bytes_received
except connection_errors:
self.failed = True
if reconnect:
await self._reconnect()
raise
if bytes_total_received < 4:
continue
elif bytes_to_receive == 0:
response_len = int.from_bytes(data[0:4], PROTOCOL_BYTE_ORDER)
bytes_to_receive = response_len
if response_len + 4 > len(data):
buffer.release()
data.extend(bytearray(response_len + 4 - len(data)))
buffer = memoryview(data)[bytes_total_received:]
continue
if bytes_total_received >= bytes_to_receive:
buffer.release()
break
buffer = buffer[bytes_received:]
return data
async def _send(self, data: Union[bytes, bytearray], reconnect=True):
if self.closed:
raise SocketError('Attempt to use closed connection.')
try:
self._writer.write(data)
await self._writer.drain()
except connection_errors:
self.failed = True
if reconnect:
await self._reconnect()
raise
def _recv(self, buffersize, flags=None) -> bytes:
"""
Handle socket data reading.
"""
kwargs = {}
if flags is not None:
kwargs['flags'] = flags
chunks = []
bytes_rcvd = 0
while bytes_rcvd < buffersize:
chunk = self.socket.recv(buffersize - bytes_rcvd, **kwargs)
if chunk == b'':
self.socket.close()
raise SocketError('Socket connection broken.')
chunks.append(chunk)
bytes_rcvd += len(chunk)
return b''.join(chunks)
def send(self, data: Union[bytes, bytearray], flags=None, reconnect=True):
"""
Send data down the socket.
:param data: bytes to send,
:param flags: (optional) OS-specific flags.
:param reconnect: (optional) reconnect on failure, default True.
"""
if self.closed:
raise SocketError('Attempt to use closed connection.')
kwargs = {}
if flags is not None:
kwargs['flags'] = flags
try:
self._socket.sendall(data, **kwargs)
except connection_errors:
self.failed = True
if reconnect:
self.reconnect()
raise
def send(self, data: bytes, flags=None):
"""
Send data down the socket.
:param data: bytes to send,
:param flags: (optional) OS-specific flags.
"""
kwargs = {}
if flags is not None:
kwargs['flags'] = flags
data = bytes(data)
total_bytes_sent = 0
while total_bytes_sent < len(data):
try:
bytes_sent = self.socket.send(data[total_bytes_sent:],
**kwargs)
except OSError:
self._socket = self.host = self.port = None
raise
if bytes_sent == 0:
self.socket.close()
raise SocketError('Socket connection broken.')
total_bytes_sent += bytes_sent