/
tcp_stream.py
65 lines (47 loc) · 1.67 KB
/
tcp_stream.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
# validated: 2017-01-07 DS 95e5295666c6 wpiutil/src/tcpsockets/TCPStream.cpp
import select
import socket
import threading
class StreamEOF(IOError):
pass
class TCPStream(object):
def __init__(self, sd, peer_ip, peer_port):
self.m_sd = sd
self.m_peerIP = peer_ip
self.m_peerPort = peer_port
self.m_rdsock = sd.makefile('rb')
self.m_wrsock = sd.makefile('wb')
self.close_lock = threading.Lock()
def read(self, size):
# TODO: ntcore does a select to wait for read to be available. Necessary?
data = self.m_rdsock.read(size)
if size > 0 and len(data) != size:
raise StreamEOF("end of file")
return data
def readStruct(self, s):
sz = s.size
data = self.m_rdsock.read(sz)
if len(data) != sz:
raise StreamEOF("end of file")
return s.unpack(data)
def send(self, contents):
self.m_wrsock.write(contents)
self.m_wrsock.flush()
def close(self):
with self.close_lock:
if self.m_sd:
try:
self.m_sd.shutdown(socket.SHUT_RDWR)
except OSError:
pass
self.m_sd.close()
self.m_sd = None
def getPeerIP(self):
return self.m_peerIP
def getPeerPort(self):
return self.m_peerPort
def setNoDelay(self):
self.m_sd.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
def _waitForReadEvent(self, timeout):
r, _, _ = select.select((self.m_sd,),(),(), timeout)
return len(r) > 0