Fix exception handling, timestamp precisions

Author: umitanuki

alpaca_trade_api/entity.py

@@ -134,6 +134,8 @@ def __getattr__(self, key):
         if key in self._raw:
             val = self._raw[key]
             if key == 'timestamp':
+                if val > 1000000000000000000:
+                    return pd.Timestamp(val, tz=NY)
                 return pd.Timestamp(val, tz=NY, unit='ms')
             return val
         return getattr(super(), key)

alpaca_trade_api/polygon/streamconn.py

@@ -98,6 +98,10 @@ async def _recv(self):
             await self.close()
             asyncio.ensure_future(self._ensure_ws())
 
+    async def consume(self):
+        if self._consume_task:
+            await self._consume_task
+
     async def _consume_msg(self):
         async for data in self._stream:
             stream = data.get('ev')

alpaca_trade_api/stream2.py

@@ -6,7 +6,7 @@
 from .common import get_base_url, get_data_url, get_credentials
 from .entity import Account, Entity, trade_mapping, agg_mapping, quote_mapping
 from . import polygon
-from .polygon.entity import Trade, Quote, Agg
+from .entity import Trade, Quote, Agg
 import logging
 
 
@@ -52,6 +52,10 @@ async def _connect(self):
 
         self._consume_task = asyncio.ensure_future(self._consume_msg())
 
+    async def consume(self):
+        if self._consume_task:
+            await self._consume_task
+
     async def _consume_msg(self):
         ws = self._ws
         try:
@@ -239,26 +243,28 @@ async def subscribe(self, channels):
     async def unsubscribe(self, channels):
         '''Handle unsubscribing from channels.'''
 
-        data_prefixes = ('Q.', 'T.', 'AM.')
-        if self._data_stream == 'polygon':
-            data_prefixes = ('Q.', 'T.', 'A.', 'AM.')
-
         data_channels = [
             c for c in channels
-            if c.startswith(data_prefixes)
+            if c.startswith(self._data_prefixes)
         ]
 
         if data_channels:
             await self.data_ws.unsubscribe(data_channels)
 
+    async def consume(self):
+        await asyncio.gather(
+            self.trading_ws.consume(),
+            self.data_ws.consume(),
+        )
+
     def run(self, initial_channels=[]):
         '''Run forever and block until exception is raised.
         initial_channels is the channels to start with.
         '''
         loop = self.loop
         try:
             loop.run_until_complete(self.subscribe(initial_channels))
-            loop.run_forever()
+            loop.run_until_complete(self.consume())
         except KeyboardInterrupt:
             logging.info("Exiting on Interrupt")
         finally: