Merge pull request #23 from alpacahq/staging

Staging Support

Author: umitanuki

.circleci/config.yml

@@ -23,6 +23,7 @@ jobs:
       - run:
           name: run tests
           command: |
+            pip install flake8 && flake8 alpaca_trade_api tests
             python setup.py test
 
       - save_cache:
@@ -58,6 +59,7 @@ jobs:
           name: run tests
           command: |
             python setup.py install
+            pip install flake8 && flake8 alpaca_trade_api tests
             python setup.py test
 
       - save_cache:

alpaca_trade_api/entity.py

@@ -74,7 +74,7 @@ def bars(self):
             t = []
             o = []
             h = []
-            l = []
+            l = []  # noqa: E741
             c = []
             v = []
             bars = []

alpaca_trade_api/polygon/__init__.py

@@ -1,2 +1,2 @@
 from .rest import REST  # noqa
-from .stream import Stream  # noqa
+from .stream import Stream  # noqa

alpaca_trade_api/polygon/entity.py

@@ -3,6 +3,7 @@
 
 NY = 'America/New_York'
 
+
 class Entity(object):
     def __init__(self, raw):
         self._raw = raw
@@ -22,7 +23,6 @@ def __repr__(self):
 
 class Agg(Entity):
     def __getattr__(self, key):
-        lkey = key.lower()
         if key in self._raw:
             val = self._raw[key]
             if key == 'day':
@@ -77,6 +77,7 @@ def df(self):
 
 class _TradeOrQuote(object):
     '''Mixin for Trade and Quote'''
+
     def __getattr__(self, key):
         if key in self._raw:
             val = self._raw[key]
@@ -116,13 +117,12 @@ def df(self):
             df.index = pd.to_datetime(
                 df.index.astype('int64') * 1000000,
                 utc=True,
-                ).tz_convert(NY)
+            ).tz_convert(NY)
 
             self._df = df
         return self._df
 
 
-
 class Trade(_TradeOrQuote, Entity):
     pass
 
@@ -148,5 +148,6 @@ class Exchange(Entity):
 class SymbolTypeMap(Entity):
     pass
 
+
 class ConditionMap(Entity):
-    pass
+    pass

alpaca_trade_api/polygon/rest.py

@@ -9,14 +9,17 @@
 
 class REST(object):
 
-    def __init__(self, api_key):
+    def __init__(self, api_key, staging=False):
         self._api_key = api_key
+        self._staging = staging
         self._session = requests.Session()
 
     def _request(self, method, path, params=None):
         url = 'https://api.polygon.io/v1' + path
         params = params or {}
         params['apiKey'] = self._api_key
+        if self._staging:
+            params['staging'] = 'true'
         resp = self._session.request(method, url, params=params)
         resp.raise_for_status()
         return resp.json()

alpaca_trade_api/polygon/stream.py

@@ -1,15 +1,12 @@
 import asyncio
 import re
 from nats.aio.client import Client as NATS
-from nats.aio.errors import ErrConnectionClosed, ErrTimeout, ErrNoServers
 
 from .entity import (
     Quote, Trade, Agg, Entity,
 )
 
-
 import json
-import time
 
 
 class Stream(object):
@@ -60,7 +57,7 @@ def _cast(self, subject, data):
                 "s": "size",
                 "t": "timestamp"
             }
-            ent = Trade({map[k]: v for k, v in data.items()})
+            ent = Trade({map[k]: v for k, v in data.items() if k in map})
         elif subject.startswith('Q.'):
             map = {
                 "sym": "symbol",
@@ -73,7 +70,7 @@ def _cast(self, subject, data):
                 "c": "condition",
                 "t": "timestamp"
             }
-            ent = Quote({map[k]: v for k, v in data.items()})
+            ent = Quote({map[k]: v for k, v in data.items() if k in map})
         elif subject.startswith('AM.'):
             map = {
                 "sym": "symbol",
@@ -90,7 +87,7 @@ def _cast(self, subject, data):
                 "e": "end",
             }
             print(data)
-            ent = Agg({map[k]: v for k, v in data.items()})
+            ent = Agg({map[k]: v for k, v in data.items() if k in map})
         else:
             ent = Entity(data)
         return ent

alpaca_trade_api/rest.py

@@ -35,7 +35,8 @@ def __init__(self, key_id=None, secret_key=None, base_url=None):
         self._session = requests.Session()
         self._retry = int(os.environ.get('APCA_MAX_RETRY', 3))
         self._retry_wait = int(os.environ.get('APCA_RETRY_WAIT', 3))
-        self.polygon = polygon.REST(self._key_id)
+        self.polygon = polygon.REST(
+            self._key_id, 'staging' in self._base_url)
 
     def _request(self, method, path, data=None, prefix='/v1'):
         url = self._base_url + prefix + path

alpaca_trade_api/stream2.py

@@ -51,7 +51,10 @@ async def _consume_msg(self):
     async def _ensure_nats(self):
         if self.polygon is not None:
             return
-        self.polygon = polygon.Stream(self._key_id)
+        key_id = self._key_id
+        if 'staging' in self._base_url:
+            key_id += '-staging'
+        self.polygon = polygon.Stream(key_id)
         self.polygon.register(r'.*', self._dispatch_nats)
         await self.polygon.connect()
 

setup.py

@@ -38,7 +38,8 @@
         'pytest-cov',
         'requests-mock',
         'coverage>=4.4.1',
-        'mock>=1.0.1'
+        'mock>=1.0.1',
+        'flake8',
     ],
     setup_requires=['pytest-runner', 'flake8'],
 )

tests/test_polygon/test_stream.py

@@ -25,6 +25,7 @@ def NATS():
 def _run(coro):
     return asyncio.get_event_loop().run_until_complete(coro)
 
+
 def test_stream(NATS):
     NATS().connect = AsyncMock()
     s = stream.Stream('api-key')
@@ -43,10 +44,9 @@ async def on_subject(stream, subject, data):
     with pytest.raises(ValueError):
         s.register('illegal', lambda x: x)
 
-
     NATS().subscribe = AsyncMock(return_value=1)
     NATS().unsubscribe = AsyncMock()
     _run(s.subscribe(['topics']))
     assert s._ssids[0] == 1
     _run(s.subscribe([]))
-    assert len(s._ssids) == 0
+    assert len(s._ssids) == 0