Added KeyboardInterrupt handling to exit cleanly

Author: bdowling

alpaca_trade_api/polygon/stream2.py

@@ -7,11 +7,12 @@
 from .entity import (
     Quote, Trade, Agg, Entity,
 )
+import logging
 
 
 class StreamConn(object):
     def __init__(self, key_id=None):
-        self._key_id = key_id
+        self._key_id = key_id or os.environ.get('APCA_API_KEY_ID')
         self._endpoint = os.environ.get(
             'POLYGON_WS_URL',
             'wss://alpaca.socket.polygon.io/stocks'
@@ -22,6 +23,7 @@ def __init__(self, key_id=None):
         self._retry = int(os.environ.get('APCA_RETRY_MAX', 3))
         self._retry_wait = int(os.environ.get('APCA_RETRY_WAIT', 3))
         self._retries = 0
+        self.loop = asyncio.get_event_loop()
 
     async def connect(self):
         await self._dispatch({'ev': 'status',
@@ -84,6 +86,9 @@ async def _recv(self):
                 msg = json.loads(r)
                 for update in msg:
                     yield update
+        except websockets.exceptions.ConnectionClosed:
+            # This error occurs when we self.close() such as on KeyboardInterrupt
+            pass
         except websockets.exceptions.ConnectionClosedError as e:
             await self._dispatch({'ev': 'status',
                                   'status': 'disconnected',
@@ -164,12 +169,15 @@ def run(self, initial_channels=[]):
         '''Run forever and block until exception is raised.
         initial_channels is the channels to start with.
         '''
-        loop = asyncio.get_event_loop()
+        loop = self.loop
         try:
             loop.run_until_complete(self.subscribe(initial_channels))
             loop.run_forever()
+        except KeyboardInterrupt:
+            logging.info("Exiting on Interrupt")
         finally:
             loop.run_until_complete(self.close())
+            loop.close()
 
     async def close(self):
         '''Close any of open connections'''

alpaca_trade_api/stream2.py

@@ -5,6 +5,7 @@
 from .common import get_base_url, get_credentials
 from .entity import Account, Entity
 from . import polygon
+import logging
 
 
 class StreamConn(object):
@@ -16,6 +17,7 @@ def __init__(self, key_id=None, secret_key=None, base_url=None):
         self._base_url = base_url
         self._ws = None
         self.polygon = None
+        self.loop = asyncio.get_event_loop()
 
     async def _connect(self):
         ws = await websockets.connect(self._endpoint)
@@ -124,12 +126,15 @@ def run(self, initial_channels=[]):
         '''Run forever and block until exception is rasised.
         initial_channels is the channels to start with.
         '''
-        loop = asyncio.get_event_loop()
+        loop = self.loop
         try:
             loop.run_until_complete(self.subscribe(initial_channels))
             loop.run_forever()
+        except KeyboardInterrupt:
+            logging.info("Exiting on Interrupt")
         finally:
             loop.run_until_complete(self.close())
+            loop.close()
 
     async def close(self):
         '''Close any of open connections'''