Add timeout to login process

itchat/components/login.py

@@ -31,7 +31,8 @@ def load_login(core):
     core.logout            = logout
 
 def login(self, enableCmdQR=False, picDir=None, qrCallback=None,
-        loginCallback=None, exitCallback=None):
+        loginCallback=None, exitCallback=None, timeout=None, timeoutCallback=None):
+    start_time = time.time()
     if self.alive or self.isLogging:
         logger.warning('itchat has already logged in.')
         return
@@ -50,6 +51,15 @@ def login(self, enableCmdQR=False, picDir=None, qrCallback=None,
             logger.info('Please scan the QR code to log in.')
         isLoggedIn = False
         while not isLoggedIn:
+            if timeout:
+                now = time.time()
+            if timeout and (now - start_time) >= timeout:
+                logger.warning('Login timeout, exit.')
+                if hasattr(timeoutCallback, '__call__'):
+                    timeoutCallback()
+                print('Timeout as user setting.')
+                return
+
             status = self.check_login()
             if hasattr(qrCallback, '__call__'):
                 qrCallback(uuid=self.uuid, status=status, qrcode=qrStorage.getvalue())

itchat/components/register.py

@@ -18,7 +18,7 @@ def load_register(core):
 
 def auto_login(self, hotReload=False, statusStorageDir='itchat.pkl',
         enableCmdQR=False, picDir=None, qrCallback=None,
-        loginCallback=None, exitCallback=None):
+        loginCallback=None, exitCallback=None, timeout=None, timeoutCallback=None):
     if not test_connect():
         logger.info("You can't get access to internet or wechat domain, so exit.")
         sys.exit()
@@ -29,11 +29,13 @@ def auto_login(self, hotReload=False, statusStorageDir='itchat.pkl',
                 loginCallback=loginCallback, exitCallback=exitCallback):
             return
         self.login(enableCmdQR=enableCmdQR, picDir=picDir, qrCallback=qrCallback,
-            loginCallback=loginCallback, exitCallback=exitCallback)
+            loginCallback=loginCallback, exitCallback=exitCallback, timeout=timeout,
+            timeoutCallback=timeoutCallback)
         self.dump_login_status(statusStorageDir)
     else:
         self.login(enableCmdQR=enableCmdQR, picDir=picDir, qrCallback=qrCallback,
-            loginCallback=loginCallback, exitCallback=exitCallback)
+            loginCallback=loginCallback, exitCallback=exitCallback, timeout=timeout,
+            timeoutCallback=timeoutCallback)
 
 def configured_reply(self):
     ''' determine the type of message and reply if its method is defined

itchat/core.py

@@ -30,7 +30,8 @@ def __init__(self):
         self.useHotReload, self.hotReloadDir = False, 'itchat.pkl'
         self.receivingRetryCount = 5
     def login(self, enableCmdQR=False, picDir=None, qrCallback=None,
-            loginCallback=None, exitCallback=None):
+            loginCallback=None, exitCallback=None, timeout=None,
+            timeoutCallback=None):
         ''' log in like web wechat does
             for log in
                 - a QR code will be downloaded and opened
@@ -45,6 +46,11 @@ def login(self, enableCmdQR=False, picDir=None, qrCallback=None,
                     - if not set, screen is cleared and qrcode is deleted
                 - exitCallback: callback after logged out
                     - it contains calling of logout
+                - timeout: the login process will forcibly end after dedicated
+                           timeout seconds
+                    - if not set, login process will always wait until user
+                      scan the QRCode to proceed.
+                - timeoutCallback: callback after timeout
             for usage
                 ..code::python
 
@@ -399,7 +405,8 @@ def load_login_status(self, fileDir,
         raise NotImplementedError()
     def auto_login(self, hotReload=False, statusStorageDir='itchat.pkl',
             enableCmdQR=False, picDir=None, qrCallback=None,
-            loginCallback=None, exitCallback=None):
+            loginCallback=None, exitCallback=None, timeout=None,
+            timeoutCallback=None):
         ''' log in like web wechat does
             for log in
                 - a QR code will be downloaded and opened
@@ -416,6 +423,11 @@ def auto_login(self, hotReload=False, statusStorageDir='itchat.pkl',
                 - exitCallback: callback after logged out
                     - it contains calling of logout
                 - qrCallback: method that should accept uuid, status, qrcode
+                - timeout: the login process will forcibly end after dedicated
+                           timeout seconds
+                    - if not set, login process will always wait until user
+                      scan the QRCode to proceed.
+                - timeoutCallback: callback after timeout
             for usage
                 ..code::python