Merge pull request #18 from vertti/refactor-analyze

Refactor analyze.py

analyze.py

@@ -13,13 +13,10 @@ logging.basicConfig(level=logging.DEBUG,
 logger = logging.getLogger(__name__)
 
 
-def get_ticker_dataframe(pair: str) -> DataFrame:
+def get_ticker(pair: str, minimum_date: arrow.Arrow) -> dict:
     """
-    Analyses the trend for the given pair
+    Request ticker data from Bittrex for a given currency pair
-    :param pair: pair as str in format BTC_ETH or BTC-ETH
-    :return: DataFrame
     """
-    minimum_date = arrow.now() - timedelta(hours=6)
     url = 'https://bittrex.com/Api/v2.0/pub/market/GetTicks'
     headers = {
         'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.36',
@@ -32,17 +29,26 @@ def get_ticker_dataframe(pair: str) -> DataFrame:
     data = requests.get(url, params=params, headers=headers).json()
     if not data['success']:
         raise RuntimeError('BITTREX: {}'.format(data['message']))
+    return data
 
-    data = [{
-        'close': t['C'],
-        'volume': t['V'],
-        'open': t['O'],
-        'high': t['H'],
-        'low': t['L'],
-        'date': t['T'],
-    } for t in sorted(data['result'], key=lambda k: k['T']) if arrow.get(t['T']) > minimum_date]
-    dataframe = DataFrame(json_normalize(data))
 
+def parse_ticker_dataframe(ticker: list, minimum_date: arrow.Arrow) -> DataFrame:
+    """
+    Analyses the trend for the given pair
+    :param pair: pair as str in format BTC_ETH or BTC-ETH
+    :return: DataFrame
+    """
+    df = DataFrame(ticker) \
+        .drop('BV', 1) \
+        .rename(columns={'C':'close', 'V':'volume', 'O':'open', 'H':'high', 'L':'low', 'T':'date'}) \
+        .sort_values('date')
+    return df[df['date'].map(arrow.get) > minimum_date]
+
+
+def populate_indicators(dataframe: DataFrame) -> DataFrame:
+    """
+    Adds several different TA indicators to the given DataFrame
+    """
     dataframe['close_30_ema'] = ta.EMA(dataframe, timeperiod=30)
     dataframe['close_90_ema'] = ta.EMA(dataframe, timeperiod=90)
 
@@ -60,37 +66,42 @@ def get_ticker_dataframe(pair: str) -> DataFrame:
     return dataframe
 
 
-def populate_trends(dataframe: DataFrame) -> DataFrame:
+def populate_buy_trend(dataframe: DataFrame) -> DataFrame:
     """
-    Populates the trends for the given dataframe
+    Based on TA indicators, populates the buy trend for the given dataframe
     :param dataframe: DataFrame
-    :return: DataFrame with populated trends
+    :return: DataFrame with buy column
-    """
     """
     dataframe.loc[
         (dataframe['stochrsi'] < 20)
-        & (dataframe['close_30_ema'] > (1 + 0.0025) * dataframe['close_60_ema']),
+        & (dataframe['macd'] > dataframe['macds'])
-        'underpriced'
-    ] = 1
-    """
-    dataframe.loc[
-        (dataframe['stochrsi'] < 20)
-        & (dataframe['macd'] > dataframe['macds']) 
         & (dataframe['close'] > dataframe['sar']),
-        'underpriced'
+        'buy'
     ] = 1
-    dataframe.loc[dataframe['underpriced'] == 1, 'buy'] = dataframe['close']
+    dataframe.loc[dataframe['buy'] == 1, 'buy_price'] = dataframe['close']
     return dataframe
 
 
+def analyze_ticker(pair: str) -> DataFrame:
+    """
+    Get ticker data for given currency pair, push it to a DataFrame and
+    add several TA indicators and buy signal to it
+    :return DataFrame with ticker data and indicator data
+    """
+    minimum_date = arrow.utcnow().shift(hours=-6)
+    data = get_ticker(pair, minimum_date)
+    dataframe = parse_ticker_dataframe(data['result'], minimum_date)
+    dataframe = populate_indicators(dataframe)
+    dataframe = populate_buy_trend(dataframe)
+    return dataframe
+
 def get_buy_signal(pair: str) -> bool:
     """
-    Calculates a buy signal based on StochRSI indicator
+    Calculates a buy signal based several technical analysis indicators
     :param pair: pair in format BTC_ANT or BTC-ANT
-    :return: True if pair is underpriced, False otherwise
+    :return: True if pair is good for buying, False otherwise
     """
-    dataframe = get_ticker_dataframe(pair)
+    dataframe = analyze_ticker(pair)
-    dataframe = populate_trends(dataframe)
     latest = dataframe.iloc[-1]
 
     # Check if dataframe is out of date
@@ -98,7 +109,7 @@ def get_buy_signal(pair: str) -> bool:
     if signal_date < arrow.now() - timedelta(minutes=10):
         return False
 
-    signal = latest['underpriced'] == 1
+    signal = latest['buy'] == 1
     logger.debug('buy_trigger: %s (pair=%s, signal=%s)', latest['date'], pair, signal)
     return signal
 
@@ -123,7 +134,7 @@ def plot_dataframe(dataframe: DataFrame, pair: str) -> None:
     ax1.plot(dataframe.index.values, dataframe['close_30_ema'], label='EMA(30)')
     ax1.plot(dataframe.index.values, dataframe['close_90_ema'], label='EMA(90)')
     # ax1.plot(dataframe.index.values, dataframe['sell'], 'ro', label='sell')
-    ax1.plot(dataframe.index.values, dataframe['buy'], 'bo', label='buy')
+    ax1.plot(dataframe.index.values, dataframe['buy_price'], 'bo', label='buy')
     ax1.legend()
 
     ax2.plot(dataframe.index.values, dataframe['macd'], label='MACD')
@@ -150,7 +161,5 @@ if __name__ == '__main__':
         pair = 'BTC_ANT'
         #for pair in ['BTC_ANT', 'BTC_ETH', 'BTC_GNT', 'BTC_ETC']:
         #   get_buy_signal(pair)
-        dataframe = get_ticker_dataframe(pair)
+        plot_dataframe(analyze_ticker(pair), pair)
-        dataframe = populate_trends(dataframe)
-        plot_dataframe(dataframe, pair)
         time.sleep(60)

test/test_analyze.py

@@ -0,0 +1,49 @@
+# pragma pylint: disable=missing-docstring
+import unittest
+from unittest.mock import patch
+from pandas import DataFrame
+import arrow
+from analyze import parse_ticker_dataframe, populate_buy_trend, populate_indicators, analyze_ticker, get_buy_signal
+
+RESULT_BITTREX = {
+    'success': True,
+    'message': '',
+    'result': [
+        {'O': 0.00065311, 'H': 0.00065311, 'L': 0.00065311, 'C': 0.00065311, 'V': 22.17210568, 'T': '2017-08-30T10:40:00', 'BV': 0.01448082},
+        {'O': 0.00066194, 'H': 0.00066195, 'L': 0.00066194, 'C': 0.00066195, 'V': 33.4727437, 'T': '2017-08-30T10:34:00', 'BV': 0.02215696},
+        {'O': 0.00065311, 'H': 0.00065311, 'L': 0.00065311, 'C': 0.00065311, 'V': 53.85127609, 'T': '2017-08-30T10:37:00', 'BV': 0.0351708},
+        {'O': 0.00066194, 'H': 0.00066194, 'L': 0.00065311, 'C': 0.00065311, 'V': 46.29210665, 'T': '2017-08-30T10:42:00', 'BV': 0.03063118},
+    ]
+}
+
+class TestAnalyze(unittest.TestCase):
+    def setUp(self):
+        self.result = parse_ticker_dataframe(RESULT_BITTREX['result'], arrow.get('2017-08-30T10:00:00'))
+
+    def test_1_dataframe_has_correct_columns(self):
+        self.assertEqual(self.result.columns.tolist(),
+                         ['close', 'high', 'low', 'open', 'date', 'volume'])
+
+    def test_2_orders_by_date(self):
+        self.assertEqual(self.result['date'].tolist(),
+                         ['2017-08-30T10:34:00',
+                          '2017-08-30T10:37:00',
+                          '2017-08-30T10:40:00',
+                          '2017-08-30T10:42:00'])
+
+    def test_3_populates_buy_trend(self):
+        dataframe = populate_buy_trend(populate_indicators(self.result))
+        self.assertTrue('buy' in dataframe.columns)
+        self.assertTrue('buy_price' in dataframe.columns)
+
+    def test_4_returns_latest_buy_signal(self):
+        buydf = DataFrame([{'buy': 1, 'date': arrow.utcnow()}])
+        with patch('analyze.analyze_ticker', return_value=buydf):
+            self.assertEqual(get_buy_signal('BTC-ETH'), True)
+        buydf = DataFrame([{'buy': 0, 'date': arrow.utcnow()}])
+        with patch('analyze.analyze_ticker', return_value=buydf):
+            self.assertEqual(get_buy_signal('BTC-ETH'), False)
+
+
+if __name__ == '__main__':
+    unittest.main()