Local client now detects when server close the connection

nemubot/__init__.py

@@ -53,12 +53,20 @@ def attach(pid, socketfile):
         sys.stderr.write("\n")
         return 1
 
-    from select import select
+    import select
+    mypoll = select.poll()
+
+    mypoll.register(sys.stdin.fileno(), select.POLLIN | select.POLLPRI)
+    mypoll.register(sock.fileno(), select.POLLIN | select.POLLPRI)
     try:
         while True:
-            rl, wl, xl = select([sys.stdin, sock], [], [])
+            for fd, flag in mypoll.poll():
+                if flag & (select.POLLERR | select.POLLHUP | select.POLLNVAL):
+                    sock.close()
+                    print("Connection closed.")
+                    return 1
 
-            if sys.stdin in rl:
+                if fd == sys.stdin.fileno():
                     line = sys.stdin.readline().strip()
                     if line == "exit" or line == "quit":
                         return 0
@@ -86,8 +94,9 @@ def attach(pid, socketfile):
                     else:
                         sock.send(line.encode() + b'\r\n')
 
-            if sock in rl:
+                if fd == sock.fileno():
                     sys.stdout.write(sock.recv(2048).decode())
+
     except KeyboardInterrupt:
         pass
     except: