Add main function to use the library and update example

Author: capaci

allokation/__init__.py

@@ -1,19 +1,8 @@
-from .functions import (
-    get_closing_price_from_yahoo,
-    map_columns_without_suffix,
-    transpose_prices,
-    calculate_multiplier,
-    calculate_amount,
-    calculate_total_for_each_ticker,
-    calculate_percentage_of_each_ticker
+from .allokate import (
+    allocate_money
 )
 
+
 __all__ = [
-    'get_closing_price_from_yahoo',
-    'map_columns_without_suffix',
-    'transpose_prices',
-    'calculate_multiplier',
-    'calculate_amount',
-    'calculate_total_for_each_ticker',
-    'calculate_percentage_of_each_ticker'
+    'allocate_money',
 ]

allokation/allokate.py

@@ -0,0 +1,25 @@
+from allokation.utils import (calculate_amount, calculate_multiplier,
+                              calculate_percentage_of_each_ticker,
+                              calculate_total_for_each_ticker,
+                              get_closing_price_from_yahoo, get_target_date,
+                              map_columns_without_suffix, transpose_prices)
+
+
+def allocate_money(available_money, tickers):
+    # import ipdb; ipdb.set_trace()
+    target_date = get_target_date()
+    prices = get_closing_price_from_yahoo(tickers=tickers, date=target_date)
+    renamed_columns = map_columns_without_suffix(tickers)
+    prices.rename(columns=renamed_columns, inplace=True)
+    df = transpose_prices(prices)
+    multiplier = calculate_multiplier(df, number_of_tickers=len(tickers), available_money=available_money)
+
+    df['amount'] = calculate_amount(df, multiplier)
+    df['total'] = calculate_total_for_each_ticker(df)
+    df['percentage'] = calculate_percentage_of_each_ticker(df)
+
+    result = {}
+    result['allocations'] = df.set_index('symbol').T.to_dict()
+    result['total_value'] = df["total"].sum()
+
+    return result

allokation/functions.py renamed to allokation/utils.py

@@ -1,7 +1,30 @@
+from datetime import date, timedelta
+
 import pandas as pd
 from pandas_datareader import data as web
 
 
+def get_target_date(base_date=date.today()):
+    weekdays = [
+        'monday',
+        'tuesday',
+        'wednesday',
+        'thursday',
+        'friday',
+        'saturday',
+        'sunday',
+    ]
+
+    target_date = base_date
+    weekday = weekdays[target_date.weekday()]
+    if weekday == 'saturday':
+        target_date = target_date - timedelta(days=1)
+    elif weekday == 'sunday':
+        target_date = target_date - timedelta(days=2)
+
+    return target_date
+
+
 def get_closing_price_from_yahoo(tickers, date):
     result = web.get_data_yahoo(tickers, date)
     return result['Adj Close']

example/example.py

@@ -1,14 +1,6 @@
-from datetime import date, timedelta
+from pprint import pp
 
-from allokation import (
-    get_closing_price_from_yahoo,
-    map_columns_without_suffix,
-    transpose_prices,
-    calculate_multiplier,
-    calculate_amount,
-    calculate_total_for_each_ticker,
-    calculate_percentage_of_each_ticker
-)
+from allokation import allocate_money
 
 AVAILABLE_MONEY = 1000
 
@@ -20,41 +12,5 @@
     'VVAR3.SA',
 ]
 
-weekdays = [
-    'monday',
-    'tuesday',
-    'wednesday',
-    'thursday',
-    'friday',
-    'saturday',
-    'sunday',
-]
-
-target_date = date.today()
-weekday = weekdays[target_date.weekday()]
-if weekday == 'saturday':
-    target_date = target_date - timedelta(days=1)
-elif weekday == 'sunday':
-    target_date = target_date - timedelta(days=2)
-
-prices = get_closing_price_from_yahoo(tickers=tickers, date=date.today()-timedelta(days=2))
-
-renamed_columns = map_columns_without_suffix(tickers)
-
-prices.rename(columns=renamed_columns, inplace=True)
-
-df = transpose_prices(prices)
-
-multiplier = calculate_multiplier(df, number_of_tickers=len(tickers), available_money=AVAILABLE_MONEY)
-
-df['amount'] = calculate_amount(df, multiplier)
-df['total'] = calculate_total_for_each_ticker(df)
-df['percentage'] = calculate_percentage_of_each_ticker(df)
-
-print(df)
-print(f'multiplier = {multiplier}')
-
-total_sum = df["total"].sum()
-print(f'MONEY = {AVAILABLE_MONEY}')
-print(f'TOTAL = {total_sum}')
-print(f'DIFF  = {AVAILABLE_MONEY - total_sum}')
+result = allocate_money(available_money=AVAILABLE_MONEY, tickers=tickers)
+pp(result)

tests/test_allokate.py

@@ -0,0 +1,41 @@
+from datetime import date
+
+import requests_cache
+from pandas_datareader import data as web
+
+from allokation.allokate import allocate_money
+
+
+def test_allocate_money(mocker):
+    """
+    The get_closing_price_from_yahoo() calls get_data_yahoo from pandas_datareader, which returns a pandas dataframe.
+    For this reason, some data for test purposes were saved in a sqlite cache.
+    In this test, the cache is used to return this data.
+    """
+    tickers = [
+        'B3SA3.SA',
+        'BBDC4.SA',
+        'MGLU3.SA',
+        'PETR4.SA',
+        'VVAR3.SA',
+    ]
+    target_date = date(year=2020, month=9, day=4)
+    cache_session = requests_cache.CachedSession(
+        cache_name='tests/data/data_from_yahoo',
+        backend='sqlite',
+        expires_after=None
+    )
+    cached_df = web.DataReader(tickers, 'yahoo', session=cache_session, start=target_date, end=target_date)
+
+    mock_result = cached_df['Adj Close']
+
+    mocker.patch('allokation.allokate.get_closing_price_from_yahoo', lambda tickers, date: mock_result.copy())
+    result = allocate_money(1000, tickers)
+
+    assert result.get('allocations', None)
+    assert result.get('total_value', None)
+    assert result['allocations'].get('B3SA3', None)
+    assert result['allocations'].get('BBDC4', None)
+    assert result['allocations'].get('MGLU3', None)
+    assert result['allocations'].get('PETR4', None)
+    assert result['allocations'].get('VVAR3', None)

tests/test_functions.py renamed to tests/test_utils.py

@@ -1,23 +1,47 @@
-import requests_cache
 import os
+from datetime import date, timedelta
+
 import pandas as pd
 import pytest
+import requests_cache
 from pandas_datareader import data as web
-from datetime import date
-
-from allokation import (
-    get_closing_price_from_yahoo,
-    map_columns_without_suffix,
-    transpose_prices,
-    calculate_multiplier,
-    calculate_amount,
-    calculate_total_for_each_ticker,
-    calculate_percentage_of_each_ticker
-)
+
+from allokation.utils import (calculate_amount, calculate_multiplier,
+                              calculate_percentage_of_each_ticker,
+                              calculate_total_for_each_ticker,
+                              get_closing_price_from_yahoo, get_target_date,
+                              map_columns_without_suffix, transpose_prices)
 
 STOCKS_DATA_FILEPATH = os.path.join(os.path.dirname(__file__), './data/stocks.csv')
 
 
+def test_get_target_date_when_today_is_a_weekday():
+    base_date = date(year=2020, month=9, day=4)
+    expected = base_date
+
+    result = get_target_date(base_date=base_date)
+
+    assert result == expected
+
+
+def test_get_target_date_when_today_is_saturday():
+    base_date = date(year=2020, month=9, day=5)
+    expected = base_date - timedelta(days=1)
+
+    result = get_target_date(base_date=base_date)
+
+    assert result == expected
+
+
+def test_get_target_date_when_today_is_sunday():
+    base_date = date(year=2020, month=9, day=6)
+    expected = base_date - timedelta(days=2)
+
+    result = get_target_date(base_date=base_date)
+
+    assert result == expected
+
+
 def test_get_closing_price_from_yahoo(mocker):
     """
     The get_closing_price_from_yahoo() calls get_data_yahoo from pandas_datareader, which returns a pandas dataframe.
@@ -45,7 +69,7 @@ def test_get_closing_price_from_yahoo(mocker):
 
     expected = cached_df['Adj Close']
 
-    mocker.patch('allokation.functions.web.get_data_yahoo', lambda tickers, date: cached_df)
+    mocker.patch('allokation.utils.web.get_data_yahoo', lambda tickers, date: cached_df)
     result = get_closing_price_from_yahoo(tickers, target_date)
 
     assert result.equals(expected)