Merge branch 'master' of https://github.com/alpacahq/alpaca-trade-api-python into feature/ws-handler-registry

Author: ttt733

README.md

@@ -30,6 +30,10 @@ account = api.get_account()
 api.list_positions()
 ```
 
+## Example Scripts
+
+Please see the `examples/` folder for some example scripts that make use of this API
+
 ## API Document
 
 The HTTP API document is located at https://docs.alpaca.markets/

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'
@@ -23,6 +24,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',
@@ -85,6 +87,9 @@ async def _recv(self):
                 msg = json.loads(r)
                 for update in msg:
                     yield update
+        except websockets.exceptions.ConnectionClosed:
+            # Ignore, occurs on self.close() such as after KeyboardInterrupt
+            pass
         except websockets.exceptions.ConnectionClosedError as e:
             await self._dispatch({'ev': 'status',
                                   'status': 'disconnected',
@@ -165,12 +170,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 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):
@@ -17,6 +18,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)
@@ -126,12 +128,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'''

examples/simpleStream.py

@@ -0,0 +1,137 @@
+#!/usr/bin/env python
+
+import os
+import sys
+import pprint
+import asyncio
+import logging
+import argparse
+
+import pandas as pd
+
+from datetime import datetime
+
+# For some fun colors...
+from colorama import Fore, Back, Style, init as ColoramaInit
+ColoramaInit(autoreset=True)
+
+# Make this global
+opt = None
+
+def ts():
+    return pd.Timestamp.now()
+
+
+def log(*args, **kwargs):
+    print(ts(), " ", *args, **kwargs)
+
+
+def debug(*args, **kwargs):
+    print(ts(), " ", *args, file=sys.stderr, **kwargs)
+
+
+def ms2date(ms, fmt='%Y-%m-%d'):
+    if isinstance(ms, pd.Timestamp):
+        return ms.strftime(fmt)
+    else:
+        return datetime.fromtimestamp(ms/1000).strftime(fmt)
+
+
+async def on_minute(conn, channel, bar):
+    symbol = bar.symbol
+    close = bar.close
+
+    try:
+        percent = (bar.dailyopen - close)/close
+        up = 1 if bar.open > bar.dailyopen else -1
+    except:  # noqa
+        percent = 0
+        up = 0
+
+    if up > 0:
+        bar_color = f'{Style.BRIGHT}{Fore.CYAN}'
+    elif up < 0:
+        bar_color = f'{Style.BRIGHT}{Fore.RED}'
+    else:
+        bar_color = f'{Style.BRIGHT}{Fore.WHITE}'
+
+    print(f'{channel:<6} {ms2date(bar.end)}  {bar_color}{symbol:<10s} {percent:>8.2f} {bar.open:>8.2f} {bar.close:>8.2f}  {bar.volume:<10d}'
+          f'  {(Fore.GREEN+"above VWAP") if close > bar.vwap else (Fore.RED+"below VWAP")}')
+
+
+async def on_tick(conn, channel, bar):
+    try:
+        percent = (bar.dailyopen - bar.close)/bar.close
+    except:  # noqa
+        percent = 0
+
+    print(f'{channel:<6s} {ms2date(bar.end)}  {bar.symbol:<10s} {percent:>8.2f}% {bar.open:>8.2f} {bar.close:>8.2f}  {bar.volume:<10d}')
+
+
+async def on_data(conn, channel, data):
+    if opt.debug or not (channel in ('AM', 'Q', 'A', 'T')):
+        debug("debug: ", pprint.pformat(data))
+
+
+
+def reloadWatch(prog, cmd):
+    async def watch_command():
+        startingmodtime = os.path.getmtime(prog)
+
+        while True:
+            modtime = os.path.getmtime(prog)
+            if modtime != startingmodtime:
+                debug(f'Reloading {" ".join(cmd)} ...')
+
+                os.execv(prog, cmd)
+
+            await asyncio.sleep(5)
+
+    return watch_command
+
+
+if __name__ == '__main__':
+    logging.basicConfig(level=logging.INFO)
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--all", "-a",
+        help="Watch the A.* feed as well, which can overwelm and backup during active times",
+        action='store_true')
+
+    parser.add_argument(
+        "--debug",
+        help="Prints debug messages",
+        action='store_true')
+
+    parser.add_argument(
+        "--polygon",
+        help="Only import Polygon.StreamConn instead of Alpaca Wrapper (mainly for testing)",
+        action='store_true')
+
+    opt = parser.parse_args()
+
+    if opt.polygon:
+        from alpaca_trade_api.polygon import StreamConn
+    else:
+        from alpaca_trade_api import StreamConn
+
+    conn = StreamConn()
+
+    # This is another way to setup wrappers for websocket callbacks, handy if conn is not global.
+    on_minute = conn.on(r'AM$')(on_minute)
+    on_tick = conn.on(r'A$')(on_tick)
+    on_data = conn.on(r'.*')(on_data)
+
+    # This is an example of how you can add your own async functions into the loop
+    # This one just watches this program for edits and tries to restart it
+    asyncio.ensure_future(reloadWatch(__file__, sys.argv)())
+
+    try:
+        if opt.all:
+            conn.run(['Q.*', 'T.*', 'AM.*', 'A.*'])
+        else:
+            conn.run(['AM.*'])
+    except Exception as e:
+        print(e)