Improve exception handling in Stream

diff --git a/tweepy/streaming.py b/tweepy/streaming.py
index 3e5173fb4889c9b7fe4d2fc7aed7582f53aad3a8..82241a3dc286ab601a459852b83ecdbc45a19c4f 100644 (file)
--- a/tweepy/streaming.py
+++ b/tweepy/streaming.py
@@ -8,7 +8,6 @@ import json
 import logging
 import re
 import ssl
-import sys
 from threading import Thread
 from time import sleep
 
@@ -217,68 +216,62 @@ class Stream:
         error_counter = 0
         resp = None
         exc_info = None
-        while self.running:
-            if self.retry_count is not None:
-                if error_counter > self.retry_count:
-                    # quit if error count greater than retry count
-                    break
-            try:
-                auth = self.auth.apply_auth()
-                resp = self.session.request('POST',
-                                            url,
-                                            data=self.body,
-                                            timeout=self.timeout,
-                                            stream=True,
-                                            auth=auth,
-                                            verify=self.verify,
-                                            proxies=self.proxies)
-                if resp.status_code != 200:
-                    if self.listener.on_error(resp.status_code) is False:
+        try:
+            while self.running:
+                if self.retry_count is not None:
+                    if error_counter > self.retry_count:
+                        # quit if error count greater than retry count
                         break
-                    error_counter += 1
-                    if resp.status_code == 420:
-                        self.retry_time = max(self.retry_420_start,
-                                              self.retry_time)
-                    sleep(self.retry_time)
-                    self.retry_time = min(self.retry_time * 2,
-                                          self.retry_time_cap)
-                else:
-                    error_counter = 0
-                    self.retry_time = self.retry_time_start
-                    self.snooze_time = self.snooze_time_step
-                    self.listener.on_connect()
-                    self._read_loop(resp)
-            except (Timeout, ssl.SSLError) as exc:
-                # This is still necessary, as a SSLError can actually be
-                # thrown when using Requests
-                # If it's not time out treat it like any other exception
-                if isinstance(exc, ssl.SSLError):
-                    if not (exc.args and 'timed out' in str(exc.args[0])):
-                        exc_info = sys.exc_info()
+                try:
+                    auth = self.auth.apply_auth()
+                    resp = self.session.request('POST',
+                                                url,
+                                                data=self.body,
+                                                timeout=self.timeout,
+                                                stream=True,
+                                                auth=auth,
+                                                verify=self.verify,
+                                                proxies=self.proxies)
+                    if resp.status_code != 200:
+                        if self.listener.on_error(resp.status_code) is False:
+                            break
+                        error_counter += 1
+                        if resp.status_code == 420:
+                            self.retry_time = max(self.retry_420_start,
+                                                  self.retry_time)
+                        sleep(self.retry_time)
+                        self.retry_time = min(self.retry_time * 2,
+                                              self.retry_time_cap)
+                    else:
+                        error_counter = 0
+                        self.retry_time = self.retry_time_start
+                        self.snooze_time = self.snooze_time_step
+                        self.listener.on_connect()
+                        self._read_loop(resp)
+                except (Timeout, ssl.SSLError) as exc:
+                    # This is still necessary, as a SSLError can actually be
+                    # thrown when using Requests
+                    # If it's not time out treat it like any other exception
+                    if isinstance(exc, ssl.SSLError):
+                        if not (exc.args and 'timed out' in str(exc.args[0])):
+                            raise
+                    if self.listener.on_timeout() is False:
                         break
-                if self.listener.on_timeout() is False:
-                    break
-                if self.running is False:
-                    break
-                sleep(self.snooze_time)
-                self.snooze_time = min(self.snooze_time + self.snooze_time_step,
-                                       self.snooze_time_cap)
-            except Exception as exc:
-                exc_info = sys.exc_info()
-                # any other exception is fatal, so kill loop
-                break
-
-        # cleanup
-        self.running = False
-        if resp:
-            resp.close()
-
-        self.new_session()
-
-        if exc_info:
-            # call a handler first so that the exception can be logged.
-            self.listener.on_exception(exc_info[1])
-            raise exc_info[1]
+                    if self.running is False:
+                        break
+                    sleep(self.snooze_time)
+                    self.snooze_time = min(
+                        self.snooze_time + self.snooze_time_step,
+                        self.snooze_time_cap
+                    )
+        except Exception as exc:
+            self.listener.on_exception(exc)
+            raise
+        finally:
+            self.running = False
+            if resp:
+                resp.close()
+            self.new_session()
 
     def _data(self, data):
         if self.listener.on_data(data) is False: