shot at simplifying sockets .. somewhat
This commit is contained in:
parent
adff0f724f
commit
f1c731bb7c
109
simplesocket.py
Normal file
109
simplesocket.py
Normal file
@ -0,0 +1,109 @@
|
||||
import errno
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
from typing import Callable, Union
|
||||
|
||||
|
||||
class Socket:
|
||||
def __init__(self, address_family: socket.AddressFamily, socket_kind: socket.SocketKind):
|
||||
self.socket = socket.socket(family=address_family, type=socket_kind)
|
||||
self.__recv_buffer = b""
|
||||
|
||||
def send(self, buffer: Union[str, bytes]):
|
||||
if isinstance(buffer, str):
|
||||
buffer = buffer.encode("UTF-8")
|
||||
|
||||
send_bytes = 0
|
||||
while send_bytes < len(buffer):
|
||||
send_bytes += self.socket.send(buffer[send_bytes:])
|
||||
|
||||
return send_bytes
|
||||
|
||||
def recv(self, maxlen: int = 4096, blocking: bool = True) -> bytes:
|
||||
maxlen -= len(self.__recv_buffer)
|
||||
try:
|
||||
if blocking:
|
||||
ret = self.__recv_buffer + self.socket.recv(maxlen)
|
||||
else:
|
||||
ret = self.__recv_buffer + self.socket.recv(maxlen, socket.MSG_DONTWAIT)
|
||||
|
||||
self.__recv_buffer = b""
|
||||
return ret
|
||||
except socket.error as e:
|
||||
err = e.args[0]
|
||||
if err == errno.EAGAIN or err == errno.EWOULDBLOCK:
|
||||
return self.__recv_buffer
|
||||
else:
|
||||
raise
|
||||
|
||||
def sendline(self, line: str):
|
||||
if not line.endswith("\n"):
|
||||
line += "\n"
|
||||
|
||||
self.send(line)
|
||||
|
||||
def recvline(self, timeout: int = 10) -> str:
|
||||
start = time.time()
|
||||
while b"\n" not in self.__recv_buffer and b"\r" not in self.__recv_buffer:
|
||||
self.__recv_buffer = self.recv(256, blocking=False)
|
||||
if time.time() - start <= timeout:
|
||||
time.sleep(0.01) # release *some* resources
|
||||
else:
|
||||
break
|
||||
|
||||
newline = b"\n"
|
||||
if newline not in self.__recv_buffer:
|
||||
newline = b"\r"
|
||||
if newline not in self.__recv_buffer:
|
||||
ret = self.__recv_buffer.decode("UTF-8")
|
||||
self.__recv_buffer = b""
|
||||
|
||||
ret = self.__recv_buffer[:self.__recv_buffer.index(newline)].decode("UTF-8")
|
||||
self.__recv_buffer = self.__recv_buffer[self.__recv_buffer.index(newline) + 1:]
|
||||
return ret
|
||||
|
||||
def close(self):
|
||||
self.socket.close()
|
||||
|
||||
|
||||
class ClientSocket(Socket):
|
||||
def __init__(self, addr: str, port: int, address_family: socket.AddressFamily = socket.AF_INET,
|
||||
socket_kind: socket.SocketKind = socket.SOCK_STREAM):
|
||||
super().__init__(address_family, socket_kind)
|
||||
self.socket.connect((addr, port))
|
||||
self.laddr, self.lport = self.socket.getsockname()
|
||||
self.raddr, self.rport = self.socket.getpeername()
|
||||
|
||||
|
||||
class RemoteSocket(Socket):
|
||||
def __init__(self, client_sock: socket.socket):
|
||||
super().__init__(client_sock.family, client_sock.type)
|
||||
self.socket = client_sock
|
||||
self.laddr, self.lport = self.socket.getsockname()
|
||||
self.raddr, self.rport = self.socket.getpeername()
|
||||
|
||||
|
||||
class ServerSocket(Socket):
|
||||
def __init__(self, addr: str, port: int, address_family: socket.AddressFamily = socket.AF_INET,
|
||||
socket_kind: socket.SocketKind = socket.SOCK_STREAM):
|
||||
super().__init__(address_family, socket_kind)
|
||||
self.socket.bind((addr, port))
|
||||
self.socket.listen(5)
|
||||
self.laddr, self.lport = self.socket.getsockname()
|
||||
self.raddr, self.rport = None, None # Transport endpoint is not connected. Surprisingly.
|
||||
|
||||
def _connection_acceptor(self, target: Callable[..., None]):
|
||||
while 1:
|
||||
(client_socket, client_address) = self.socket.accept()
|
||||
connection_handler_thread = threading.Thread(target=target, args=(RemoteSocket(client_socket), ))
|
||||
connection_handler_thread.start()
|
||||
|
||||
def accept(self, target: Callable[..., None], blocking: bool = True):
|
||||
if blocking:
|
||||
self._connection_acceptor(target)
|
||||
return None
|
||||
else:
|
||||
connection_accept_thread = threading.Thread(target=self._connection_acceptor, kwargs={'target': target})
|
||||
connection_accept_thread.start()
|
||||
return connection_accept_thread
|
||||
Loading…
Reference in New Issue
Block a user