micropython · devxpy · Jun 21, 2018
diff --git a/umqtt.simple/umqtt/simple.py b/umqtt.simple/umqtt/simple.py
@@ -2,13 +2,24 @@
 import ustruct as struct
 from ubinascii import hexlify
 
+
 class MQTTException(Exception):
     pass
 
-class MQTTClient:
 
-    def __init__(self, client_id, server, port=0, user=None, password=None, keepalive=0,
-                 ssl=False, ssl_params={}):
+class MQTTClient:
+    def __init__(
+        self,
+        client_id,
+        server,
+        port=0,
+        user=None,
+        password=None,
+        keepalive=0,
+        ssl=False,
+        ssl_params={},
+        timeout=None,
+    ):
         if port == 0:
             port = 8883 if ssl else 1883
         self.client_id = client_id
@@ -26,6 +37,7 @@ def __init__(self, client_id, server, port=0, user=None, password=None, keepaliv
         self.lw_msg = None
         self.lw_qos = 0
         self.lw_retain = False
+        self.timeout = timeout
 
     def _send_str(self, s):
         self.sock.write(struct.pack("!H", len(s)))
@@ -58,7 +70,10 @@ def connect(self, clean_session=True):
         self.sock.connect(addr)
         if self.ssl:
             import ussl
+
             self.sock = ussl.wrap_socket(self.sock, **self.ssl_params)
+        self.sock.settimeout(self.timeout)
+
         premsg = bytearray(b"\x10\0\0\0\0\0")
         msg = bytearray(b"\x04MQTT\x04\x02\0\0")
 
@@ -85,7 +100,7 @@ def connect(self, clean_session=True):
 
         self.sock.write(premsg, i + 2)
         self.sock.write(msg)
-        #print(hex(len(msg)), hexlify(msg, ":"))
+        # print(hex(len(msg)), hexlify(msg, ":"))
         self._send_str(self.client_id)
         if self.lw_topic:
             self._send_str(self.lw_topic)
@@ -119,7 +134,7 @@ def publish(self, topic, msg, retain=False, qos=0):
             sz >>= 7
             i += 1
         pkt[i] = sz
-        #print(hex(len(pkt)), hexlify(pkt, ":"))
+        # print(hex(len(pkt)), hexlify(pkt, ":"))
         self.sock.write(pkt, i + 1)
         self._send_str(topic)
         if qos > 0:
@@ -146,15 +161,15 @@ def subscribe(self, topic, qos=0):
         pkt = bytearray(b"\x82\0\0\0")
         self.pid += 1
         struct.pack_into("!BH", pkt, 1, 2 + 2 + len(topic) + 1, self.pid)
-        #print(hex(len(pkt)), hexlify(pkt, ":"))
+        # print(hex(len(pkt)), hexlify(pkt, ":"))
         self.sock.write(pkt)
         self._send_str(topic)
         self.sock.write(qos.to_bytes(1, "little"))
         while 1:
             op = self.wait_msg()
             if op == 0x90:
                 resp = self.sock.read(4)
-                #print(resp)
+                # print(resp)
                 assert resp[1] == pkt[2] and resp[2] == pkt[3]
                 if resp[3] == 0x80:
                     raise MQTTException(resp[3])
@@ -166,7 +181,7 @@ def subscribe(self, topic, qos=0):
     # messages processed internally.
     def wait_msg(self):
         res = self.sock.read(1)
-        self.sock.setblocking(True)
+        # self.sock.setblocking(True)
         if res is None:
             return None
         if res == b"":
@@ -200,5 +215,8 @@ def wait_msg(self):
     # If not, returns immediately with None. Otherwise, does
     # the same processing as wait_msg.
     def check_msg(self):
-        self.sock.setblocking(False)
-        return self.wait_msg()
+        try:
+            self.sock.setblocking(False)
+            return self.wait_msg()
+        finally:
+            self.sock.settimeout(self.timeout)