1import asyncutils as A, asyncio as I
2from sys import audit
3from asyncutils._internal.py312 import Queue
4from asyncutils._internal.helpers import LoopMixinBase, fullname
5from asyncutils._internal.log import warning
6from asyncutils._internal.submodules import networking_all as __all__
[docs]
7class LineProtocol(I.Protocol, LoopMixinBase):
8 NEWLINE, CARRIAGE_RETURN, _h = __import__('os').linesep.encode(), b'\r', A.ignore_cancellation.combined(I.InvalidStateError); __slots__ = '__buf', '__cl', '__dw', '__er', '__lines', '__paused', 'transport'
9 def __init__(self): audit(fullname(self)); self.__buf, self.__lines = bytearray(), Queue(); self.__cl = self.__paused = self.__er = False; self.transport = self.__dw = None
10 @property
11 def connected_transport(self):
12 if (t := self.transport) is None: raise ConnectionError('asyncutils.networking.LineProtocol: no transport connected')
13 return t
[docs]
14 def connection_made(self, transport): self.transport = transport
[docs]
15 def connection_lost(self, exc):
16 if t := self.transport:
17 with self._h: t.abort()
18 self.__lines.shutdown(); self.__cl = True
19 if w := self.__dw:
20 if not w.done(): w.set_exception(ConnectionError('asyncutils.networking.LineProtocol: transport connection lost') if exc is None else exc)
21 self.__dw = None
[docs]
22 def close(self):
23 if t := self.transport:
24 with self._h: t.close(); self.__cl = True
25 return self.__cl
[docs]
26 def data_received(self, data, bufsize=None):
27 if bufsize is None: bufsize = A.getcontext().LINE_PROTOCOL_DEFAULT_BUFFER_SIZE
28 (b := self.__buf).extend(data); n = self.NEWLINE
29 if len(b) > bufsize: self.flush()
30 while not self.__cl and n in b: l, b = b.split(n, 1); self.__pl(l)
31 self.__buf = b
32 if self.__er: self.flush(); self.signal_eof()
[docs]
33 def flush(self): self.__pl(b := self.__buf); b.clear()
[docs]
34 def signal_eof(self): self.__lines.put_nowait(None)
[docs]
35 def pause_writing(self):
36 self.__paused = True
37 if self.__dw is None: self.__dw = self.make_fut()
[docs]
38 def resume_writing(self):
39 self.__paused = False
40 if w := self.__dw:
41 if not w.done(): w.set_result(None)
42 self.__dw = None
43 def __pl(self, data): self.__lines.put_nowait(data.rstrip(self.CARRIAGE_RETURN).decode('utf-8'))
[docs]
44 def write_line(self, line):
45 with self._h:
46 if not self.connected_transport.is_closing(): self.write_literal(line.encode('utf-8', 'ignore')+self.NEWLINE)
[docs]
47 def write_literal(self, data): self.connected_transport.write(data)
[docs]
48 def eof_received(self):
49 self.__er = True
50 if self.__buf: self.flush()
51 if self.__lines.empty() and not self.__cl: self.signal_eof()
[docs]
52 async def read_line(self): L = self.__lines; return None if self.__cl and L.empty() else (L.task_done() if (l := await L.get()) is None else l)
[docs]
53 async def drain(self):
54 if self.__paused and (w := self.__dw): await w
[docs]
55 async def write_line_with_backpressure(self, line): await self.drain(); self.write_line(line)
[docs]
56 async def write_literal_with_backpressure(self, data): await self.drain(); self.write_literal(data)
[docs]
57class LFProtocol(LineProtocol): NEWLINE, __slots__ = b'\n', ()
[docs]
58class CRLFProtocol(LineProtocol): NEWLINE, __slots__ = b'\r\n', ()
[docs]
59class CRProtocol(LineProtocol): NEWLINE, __slots__ = b'\r', ()
[docs]
60class SocketTransport(I.Transport):
61 __slots__ = '__buf', '__cl', '__lim', '__protocol', '__sock'; _h, ptc = A.IgnoreErrors(OSError), LineProtocol
62 @property
63 def loop(self): return self.__protocol.loop
64 def __init__(self, sock=None):
65 audit(fullname(self)); self.__rx(); (p := self.ptc()).connection_made(self); self.__sock, self.__cl, self.__buf, self.__lim, self.__protocol = sock, False, bytearray(), A.getcontext().SOCKET_TRANSPORT_LIMITS, p
66 if sock is not None: self.connect_sock(sock)
67 def __rx(self, _=('socket', 'sockname', 'peername')): super().__init__(dict.fromkeys(_))
68 def __rr(self, sock, size=None):
69 try: self.__protocol.data_received(d) if (d := sock.recv(A.getcontext().LINE_PROTOCOL_DEFAULT_BUFFER_SIZE if size is None else size)) else (self.__protocol.eof_received() or self.close())
70 except OSError as e: warning('%s: read error', fullname(self)); self.close(e)
[docs]
71 def connect_sock(self, sock=None):
72 if sock is None and (sock := self.__sock) is None: return
73 sock.setblocking(False); self.loop.add_reader(sock.fileno(), self.__rr, sock); (e := self._extra)['sockname'] = sock.getsockname() # ty: ignore[unresolved-attribute]
74 with self._h: e['peername'] = sock.getpeername()
[docs]
75 def disconnect_sock(self):
76 if (s := self.__sock) is None: return s
77 with self._h: s.close()
78 self.loop.remove_reader(s.fileno()); self.__sock = None; self.__rx(); return s
[docs]
79 @A.dualcontextmanager
80 def sock_context(self, sock):
81 try: yield self.connect_sock(sock)
82 finally: self.disconnect_sock()
83 def __writer(self, data, bufsize=None):
84 if self.__cl: return
85 (b := self.__buf).extend(data)
86 if bufsize is None: bufsize = A.getcontext().SOCKET_TRANSPORT_LIMITS[1]
87 if len(b) <= bufsize or (s := self.__sock) is None: return
88 try: s.sendall(b); b.clear()
89 except OSError as e: warning('%s: write error', fullname(self)); self.close(e)
[docs]
90 def write(self, data): self.loop.call_soon(self.__writer, data)
[docs]
91 def get_write_buffer_size(self): return len(self.__buf)
[docs]
92 def get_write_buffer_limits(self): return self.__lim
[docs]
93 def set_write_buffer_limits(self, high=None, low=None):
94 l = self.__lim
95 if low is None: low = l[0]
96 if high is None: high = l[1]
97 self.__lim = min(max(low, 0), high := min(high, A.getcontext().SOCKET_TRANSPORT_LIMITS[1])), high
[docs]
98 def write_eof(self):
99 if not (self.__cl or (s := self.__sock) is None):
100 with self._h: s.shutdown(1)
[docs]
101 def can_write_eof(self): return True # noqa: PLR6301
[docs]
102 def is_closing(self): return self.__cl
[docs]
103 def close(self, e=None):
104 if self.__cl: return
105 self.__cl = True; self.__protocol.connection_lost(e); self.disconnect_sock()
[docs]
106 def get_protocol(self): return self.__protocol
107 def set_protocol(self, protocol):
108 if not isinstance(protocol, LineProtocol): raise TypeError('asyncutils.networking.SocketTransport: protocol should be a LineProtocol')
109 self.__protocol.connection_lost(None); protocol.connection_made(self); self.__protocol = protocol; self.connect_sock()
[docs]
110 def abort(self): self.loop.call_soon(self.close)