Add a way to define the percentages wanted for each stock

README.md

@@ -23,7 +23,7 @@ pip install allokation
 
 ## Usage
 
-- It's quite simple to use this package, you just need to import the function `allocate_money`, pass a list of tickers you want and the available money you have to invest.
+- It's quite simple to use this package, you just need to import the function `allocate_money`, pass a list of tickers you want and the available money you have to invest. If you want, you can also pass a list of the percentages of each stocks you want in your portfolio. This list of percentages must have the same length of the tickers.
 
 - It will return a dict containing the allocations you must need and the total money you must need to have this portfolio (This total will be less or equal than the available money you informed to the `allocate_money` function). For each stock, it will be returned the `price` that was used to calculate the portfolio, the `amount` of stocks you will need to buy, the `total` money you need to buy this amount of this stock and the `percentage` that this stock represents in your portfolio. For example:
 

allokation/allokate.py

@@ -1,25 +1,28 @@
-from allokation.utils import (calculate_amount, calculate_multiplier,
+from allokation.utils import (calculate_amount,
                               calculate_percentage_of_each_ticker,
                               calculate_total_for_each_ticker,
-                              get_closing_price_from_yahoo, get_target_date,
-                              map_columns_without_suffix, transpose_prices)
+                              get_closing_price_from_yahoo,
+                              get_percentage_of_stocks, get_target_date,
+                              transpose_prices)
 
 
-def allocate_money(available_money, tickers):
-    # import ipdb; ipdb.set_trace()
+def allocate_money(available_money, tickers, percentages=None):
+    if percentages and len(tickers) != len(percentages):
+        raise Exception('Tickers and percentages must have the same lenght')
+
+    percentage_multiplier = get_percentage_of_stocks(tickers=tickers, percentages=percentages)
+
     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['amount'] = calculate_amount(df, available_money, percentage_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()
+    result['total_value'] = df["total"].sum().round(2)
 
     return result

allokation/utils.py

@@ -4,6 +4,13 @@
 from pandas_datareader import data as web
 
 
+def get_percentage_of_stocks(tickers, percentages=None):
+    if percentages:
+        return pd.Series(percentages)/100
+
+    return 1/(len(tickers))
+
+
 def get_target_date(base_date=date.today()):
     weekdays = [
         'monday',
@@ -30,10 +37,6 @@ def get_closing_price_from_yahoo(tickers, date):
     return result['Adj Close']
 
 
-def map_columns_without_suffix(tickers):
-    return {ticker: ticker[:-3] for ticker in tickers}
-
-
 def transpose_prices(prices):
     df = pd.DataFrame()
     df['symbol'] = prices.columns
@@ -42,15 +45,8 @@ def transpose_prices(prices):
     return df
 
 
-def calculate_multiplier(df, number_of_tickers, available_money):
-    max_price = df['price'].max()
-    percentage = 1 / number_of_tickers
-    multiplier = (available_money * percentage) / max_price
-    return multiplier
-
-
-def calculate_amount(df, multiplier):
-    return (df['price'].max()*multiplier/df['price']).round(0)
+def calculate_amount(df, available_money, percentage_multiplier):
+    return (available_money * percentage_multiplier / df['price']).round(0)
 
 
 def calculate_total_for_each_ticker(df):

example/example.py

@@ -12,5 +12,7 @@
     'VVAR3.SA',
 ]
 
-result = allocate_money(available_money=AVAILABLE_MONEY, tickers=tickers)
+percentages = [60, 10, 10, 10, 10]
+
+result = allocate_money(available_money=AVAILABLE_MONEY, tickers=tickers, percentages=percentages)
 pp(result)

tests/test_allokate.py

@@ -1,5 +1,6 @@
 from datetime import date
 
+import pytest
 import requests_cache
 from pandas_datareader import data as web
 
@@ -34,8 +35,23 @@ def test_allocate_money(mocker):
 
     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)
+    assert result['allocations'].get('B3SA3.SA', None)
+    assert result['allocations'].get('BBDC4.SA', None)
+    assert result['allocations'].get('MGLU3.SA', None)
+    assert result['allocations'].get('PETR4.SA', None)
+    assert result['allocations'].get('VVAR3.SA', None)
+
+
+def test_allocate_money_with_percentage_with_different_length_than_stocks():
+    with pytest.raises(Exception):
+        tickers = [
+            'B3SA3.SA',
+            'BBDC4.SA',
+            'MGLU3.SA',
+            'PETR4.SA',
+            'VVAR3.SA',
+        ]
+
+        percentages = [10, 10, 10, 10, 10, 10, 10, 10, 10, 10]
+
+        allocate_money(1000, tickers, percentages=percentages)

tests/test_utils.py

@@ -2,19 +2,49 @@
 from datetime import date, timedelta
 
 import pandas as pd
-import pytest
 import requests_cache
 from pandas_datareader import data as web
 
-from allokation.utils import (calculate_amount, calculate_multiplier,
+from allokation.utils import (calculate_amount,
                               calculate_percentage_of_each_ticker,
                               calculate_total_for_each_ticker,
-                              get_closing_price_from_yahoo, get_target_date,
-                              map_columns_without_suffix, transpose_prices)
+                              get_closing_price_from_yahoo,
+                              get_percentage_of_stocks, get_target_date,
+                              transpose_prices)
 
 STOCKS_DATA_FILEPATH = os.path.join(os.path.dirname(__file__), './data/stocks.csv')
 
 
+def test_get_percentage_of_stocks_without_percentage_should_return_equal_distribution():
+    tickers = [
+        'B3SA3.SA',
+        'BBDC4.SA',
+        'CSAN3.SA',
+        'CYRE3.SA',
+    ]
+
+    expected = 0.25
+
+    result = get_percentage_of_stocks(tickers=tickers)
+
+    assert result == expected
+
+
+def test_get_percentage_of_stocks_with_percentage_should_return_pandas_series():
+    tickers = [
+        'B3SA3.SA',
+        'BBDC4.SA',
+        'CSAN3.SA',
+        'CYRE3.SA',
+    ]
+    percentages = [40, 20, 20, 20]
+    expected = pd.Series([0.4, 0.2, 0.2, 0.2])
+
+    result = get_percentage_of_stocks(tickers=tickers, percentages=percentages)
+
+    assert result.equals(expected)
+
+
 def test_get_target_date_when_today_is_a_weekday():
     base_date = date(year=2020, month=9, day=4)
     expected = base_date
@@ -88,42 +118,26 @@ def test_transpose_prices():
     assert result.equals(expected)
 
 
-def test_map_columns_without_suffix():
-    tickers = [
-        'MGLU3.SA',
-        'PETR4.SA',
-        'VVAR3.SA',
-    ]
-
-    result = map_columns_without_suffix(tickers)
-
-    expected = {
-        'MGLU3.SA': 'MGLU3',
-        'PETR4.SA': 'PETR4',
-        'VVAR3.SA': 'VVAR3',
-    }
-    assert result == expected
-
-
-def test_calculate_multiplier():
+def test_calculate_amount_with_equal_distribution():
     df = pd.read_csv(STOCKS_DATA_FILEPATH)
-    number_of_tickers = len(df.values)
     available_money = 1000
+    percentage_multiplier = 1/len(df)
 
-    expected = pytest.approx(3.57, 0.01)
+    expected = (available_money*percentage_multiplier/df['price']).round(0)
 
-    result = calculate_multiplier(df, number_of_tickers, available_money)
+    result = calculate_amount(df, available_money=available_money, percentage_multiplier=percentage_multiplier)
 
-    assert result == expected
+    assert result.equals(expected)
 
 
 def test_calculate_amount():
     df = pd.read_csv(STOCKS_DATA_FILEPATH)
-    multiplier = 1
+    available_money = 1000
+    percentage_multiplier = pd.Series([0.33, 0.33, 0.34])
 
-    expected = (df['price'].max()/df['price']).round(0)
+    expected = (available_money*percentage_multiplier/df['price']).round(0)
 
-    result = calculate_amount(df, multiplier=multiplier)
+    result = calculate_amount(df, available_money=available_money, percentage_multiplier=percentage_multiplier)
 
     assert result.equals(expected)