From 6fc85653cb27ef42ee4d344636651dced168492b Mon Sep 17 00:00:00 2001
From: bspeice <bspeice.nc@gmail.com>
Date: Thu, 24 Nov 2016 18:27:29 -0500
Subject: [PATCH] Add fetching options quotes (daily data from Tradeking)

---
 metrik/tasks/tradeking.py    | 38 +++++++++++++++++++++++++++++++++++-
 test/tasks/test_tradeking.py | 17 ++++++++++++++--
 2 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/metrik/tasks/tradeking.py b/metrik/tasks/tradeking.py
index d704813..ba361a4 100644
--- a/metrik/tasks/tradeking.py
+++ b/metrik/tasks/tradeking.py
@@ -9,6 +9,13 @@ from metrik.tasks.base import MongoCreateTask, MongoRateLimit, MongoNoBackCreate
 from metrik.conf import get_config
 
 
+def batch(iterable, size):
+    sourceiter = iter(iterable)
+    while True:
+        batchiter = islice(sourceiter, size)
+        yield chain([batchiter.next()], batchiter)
+
+
 class TradekingApi(object):
     format_json = '.json'
     format_xml = '.xml'
@@ -87,6 +94,7 @@ class Tradeking1mTimesales(MongoCreateTask):
 
 
 class TradekingOptionsQuotes(MongoNoBackCreateTask):
+    batch_size = 75
     symbol = Parameter()
 
     def get_collection_name(self):
@@ -101,9 +109,19 @@ class TradekingOptionsQuotes(MongoNoBackCreateTask):
 
         return [r['symbol'] for r in results]
 
+    @staticmethod
+    def retrieve_quotes(api, symbols):
+        response = api.api_request('market/ext/quotes', {
+            'symbols': ','.join(symbols)
+        })
+
+        results = response.json()['response']['quotes']['quote']
+
+        return results
+
     @staticmethod
     def retrieve_data(symbol):
-        tradeking = TradekingApi()
+        api = TradekingApi()
 
         # We request a first rate limit lock to get the options chain
         ratelimit = MongoRateLimit()
@@ -112,3 +130,21 @@ class TradekingOptionsQuotes(MongoNoBackCreateTask):
             limit=60,
             interval=timedelta(minutes=1)
         )
+
+        if not chain_acquire:
+            return {}
+
+
+        chain = TradekingOptionsQuotes.retrieve_chain_syms(api, symbol)
+        results = []
+        for b in batch(chain, TradekingOptionsQuotes.batch_size):
+            batch_acquire = ratelimit.acquire_lock(
+                service='tradeking',
+                limit=60,
+                interval=timedelta(minutes=1)
+            )
+            if batch_acquire:
+                batch_results = TradekingOptionsQuotes.retrieve_quotes(api, b)
+                results += batch_results
+
+        return results
diff --git a/test/tasks/test_tradeking.py b/test/tasks/test_tradeking.py
index 15641d9..962d0de 100644
--- a/test/tasks/test_tradeking.py
+++ b/test/tasks/test_tradeking.py
@@ -6,7 +6,7 @@ import pytest
 from six.moves import map
 from pytz import utc
 
-from metrik.tasks.tradeking import Tradeking1mTimesales, TradekingOptionsQuotes, TradekingApi
+from metrik.tasks.tradeking import Tradeking1mTimesales, TradekingOptionsQuotes, TradekingApi, batch
 from metrik.trading_days import TradingDay
 from metrik.targets.mongo import MongoTarget
 from test.mongo_test import MongoTest
@@ -40,6 +40,7 @@ def test_returns_verifiable(ticker):
 
     assert_allclose(tradeking_ohlc, yahoo_ohlc, rtol=1e-3)
 
+
 @pytest.mark.parametrize('ticker', [
     'AAPL', 'GOOG', 'SPY', 'REGN', 'SWHC', 'BAC', 'NVCR', 'ARGT'
 ])
@@ -49,8 +50,20 @@ def test_chain_returns(ticker):
 
     assert len(chain) > 20
 
-class TradekingTest(MongoTest):
 
+@pytest.mark.parametrize('ticker', [
+    'AAPL', 'GOOG', 'SPY'
+])
+def test_options_quote(ticker):
+    demo_length = 25
+    api = TradekingApi()
+    chain = TradekingOptionsQuotes.retrieve_chain_syms(api, ticker)[:demo_length]
+    quotes = TradekingOptionsQuotes.retrieve_quotes(api, chain)
+
+    assert len(quotes) == demo_length
+
+
+class TradekingTest(MongoTest):
     def test_record_is_saveable(self):
         # Had an issue previously where the `date` type would cause saving
         # to fail. Make sure that doesn't continue