Merge pull request #461 from plotly/https_stream

Https stream

Author: BRONSOLO

CHANGELOG.md

@@ -4,6 +4,10 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 
 ## [Unreleased]
 
+## [1.9.12] - 2016-05-16
+### Added
+SSL support for streaming.
+
 ## [1.9.11] - 2016-05-02
 ### Added
 - The FigureFactory can now create scatter plot matrices with `.create_scatterplotmatrix`. Check it out with:

plotly/plotly/chunked_requests/chunked_request.py

@@ -3,10 +3,12 @@
 import os
 from six.moves import http_client
 from six.moves.urllib.parse import urlparse
+from ssl import SSLError
+
 
 class Stream:
-    def __init__(self, server, port=80, headers={}, url='/'):
-        ''' Initialize a stream object and an HTTP Connection
+    def __init__(self, server, port=80, headers={}, url='/', ssl_enabled=False):
+        ''' Initialize a stream object and an HTTP or HTTPS connection
         with chunked Transfer-Encoding to server:port with optional headers.
         '''
         self.maxtries = 5
@@ -17,6 +19,7 @@ def __init__(self, server, port=80, headers={}, url='/'):
         self._port = port
         self._headers = headers
         self._url = url
+        self._ssl_enabled = ssl_enabled
         self._connect()
 
     def write(self, data, reconnect_on=('', 200, )):
@@ -73,17 +76,21 @@ def write(self, data, reconnect_on=('', 200, )):
     def _get_proxy_config(self):
         """
         Determine if self._url should be passed through a proxy. If so, return
-        the appropriate proxy_server and proxy_port
+        the appropriate proxy_server and proxy_port. Assumes https_proxy is used
+        when ssl_enabled=True.
 
         """
 
         proxy_server = None
         proxy_port = None
+        ssl_enabled = self._ssl_enabled
 
-        ## only doing HTTPConnection, so only use http_proxy
-        proxy = os.environ.get("http_proxy")
+        if ssl_enabled:
+            proxy = os.environ.get("https_proxy")
+        else:
+            proxy = os.environ.get("http_proxy")
         no_proxy = os.environ.get("no_proxy")
-        no_proxy_url = no_proxy and self._url in no_proxy
+        no_proxy_url = no_proxy and self._server in no_proxy
 
         if proxy and not no_proxy_url:
             p = urlparse(proxy)
@@ -93,19 +100,30 @@ def _get_proxy_config(self):
         return proxy_server, proxy_port
 
     def _connect(self):
-        ''' Initialize an HTTP connection with chunked Transfer-Encoding
+        ''' Initialize an HTTP/HTTPS connection with chunked Transfer-Encoding
         to server:port with optional headers.
         '''
         server = self._server
         port = self._port
         headers = self._headers
+        ssl_enabled = self._ssl_enabled
         proxy_server, proxy_port = self._get_proxy_config()
 
         if (proxy_server and proxy_port):
-            self._conn = http_client.HTTPConnection(proxy_server, proxy_port)
+            if ssl_enabled:
+                self._conn = http_client.HTTPSConnection(
+                    proxy_server, proxy_port
+                )
+            else:
+                self._conn = http_client.HTTPConnection(
+                    proxy_server, proxy_port
+                )
             self._conn.set_tunnel(server, port)
         else:
-            self._conn = http_client.HTTPConnection(server, port)
+            if ssl_enabled:
+                self._conn = http_client.HTTPSConnection(server, port)
+            else:
+                self._conn = http_client.HTTPConnection(server, port)
 
         self._conn.putrequest('POST', self._url)
         self._conn.putheader('Transfer-Encoding', 'chunked')
@@ -236,6 +254,16 @@ def _isconnected(self):
                 # let's just assume that we're still connected and
                 # hopefully recieve some data on the next try.
                 return True
+            elif isinstance(e, SSLError):
+                if e.errno == 2:
+                    # errno 2 occurs when trying to read or write data, but more
+                    # data needs to be received on the underlying TCP transport
+                    # before the request can be fulfilled.
+                    #
+                    # Python 2.7.9+ and Python 3.3+ give this its own exception,
+                    # SSLWantReadError
+                    return True
+                raise e
             else:
                 # Unknown scenario
                 raise e

plotly/plotly/plotly.py

@@ -439,6 +439,9 @@ class Stream:
 
     """
 
+    HTTP_PORT = 80
+    HTTPS_PORT = 443
+
     @utils.template_doc(**tools.get_config_file())
     def __init__(self, stream_id):
         """
@@ -454,6 +457,30 @@ def __init__(self, stream_id):
         self.connected = False
         self._stream = None
 
+    def get_streaming_specs(self):
+        """
+        Returns the streaming server, port, ssl_enabled flag, and headers.
+
+        """
+        streaming_url = get_config()['plotly_streaming_domain']
+        ssl_enabled = 'https' in streaming_url
+        port = self.HTTPS_PORT if ssl_enabled else self.HTTP_PORT
+
+        # If no scheme (https/https) is included in the streaming_url, the
+        # host will be None. Use streaming_url in this case.
+        host = (six.moves.urllib.parse.urlparse(streaming_url).hostname or
+                streaming_url)
+
+        headers = {'Host': host, 'plotly-streamtoken': self.stream_id}
+        streaming_specs = {
+            'server': host,
+            'port': port,
+            'ssl_enabled': ssl_enabled,
+            'headers': headers
+        }
+
+        return streaming_specs
+
     def heartbeat(self, reconnect_on=(200, '', 408)):
         """
         Keep stream alive. Streams will close after ~1 min of inactivity.
@@ -481,10 +508,8 @@ def open(self):
         https://plot.ly/python/streaming/
 
         """
-        streaming_url = get_config()['plotly_streaming_domain']
-        self._stream = chunked_requests.Stream(
-            streaming_url, 80, {'Host': streaming_url,
-                                'plotly-streamtoken': self.stream_id})
+        streaming_specs = self.get_streaming_specs()
+        self._stream = chunked_requests.Stream(**streaming_specs)
 
     def write(self, trace, layout=None, validate=True,
               reconnect_on=(200, '', 408)):