diff --git a/flat2mongo.py b/flat2mongo.py new file mode 100644 index 0000000..125b7f9 --- /dev/null +++ b/flat2mongo.py @@ -0,0 +1,26 @@ +# Utility script to port from flatfile to database + +import config as cfg +import flatfile as f +import mongo as m +import json + +with open(cfg.database) as f: + data = json.loads(f.read()) + +for user in data.keys(): + print("\n\nUser:", user) + print("Portfolio:") + + for stock in data[user]["portfolio"].keys(): + ticker = stock + amount = data[user]["portfolio"][stock] + m.add_stock(user, ticker, amount) + + print("Watchlist:") + for stock in data[user]["watchlist"].keys(): + ticker = stock + amount = data[user]["watchlist"][stock] + m.watch(user, ticker, amount) + +print("[*] Done!") diff --git a/main.py b/main.py index 644500c..f4b397a 100644 --- a/main.py +++ b/main.py @@ -6,7 +6,6 @@ import config as cfg from decimal import Decimal import typing -print(cfg.use_mongodb) if cfg.use_mongodb: import mongo as database diff --git a/mongo.py b/mongo.py index cbd623d..5a72c3a 100644 --- a/mongo.py +++ b/mongo.py @@ -17,6 +17,14 @@ class Watchlist(Document): stocks = DictField(default={}) +def encodeKey(key): + return key.replace(".", "\\u002e") + + +def decodeKey(key): + return key.replace("\\u002e", ".") + + def user_check(user): # Will create empty portfolio and watchlist for new users if not Portfolio.objects(user=user): @@ -32,7 +40,11 @@ def user_check(user): def get_stocks(user): user_check(user) - return Portfolio.objects.get(user=user).stocks + stocks = Portfolio.objects.get(user=user).stocks + decoded_stocks = {} + for stock, amount in stocks.items(): + decoded_stocks[decodeKey(stock)] = amount.to_decimal() + return decoded_stocks def add_stock(user, stock, amount): @@ -43,7 +55,7 @@ def add_stock(user, stock, amount): if not (yfi.stock_exists(stock)): # Stock ticker does not exist on yahoo finance return False - + stock = encodeKey(stock) if stock in portfolio.stocks.keys(): portfolio.stocks[stock] = Decimal128( amount + portfolio.stocks[stock].to_decimal() @@ -56,6 +68,7 @@ def add_stock(user, stock, amount): def delete_stock(user, stock): user_check(user) + stock = encodeKey(stock) portfolio = Portfolio.objects.get(user=user) if stock in portfolio.stocks.keys(): del portfolio.stocks[stock] @@ -67,24 +80,30 @@ def delete_stock(user, stock): def get_watchlist(user): user_check(user) - return Watchlist.objects.get(user=user).stocks + wlist = Watchlist.objects.get(user=user).stocks + decoded_stocks = {} + for stock, val in wlist.items(): + decoded_stocks[decodeKey(stock)] = val.to_decimal() + return decoded_stocks def watch(user, stock, est_price="0"): user_check(user) + watchlist = Watchlist.objects.get(user=user) if not (yfi.stock_exists(stock)): # Stock ticker does not exist on yahoo finance return False - - watchlist.stocks[stock] = str(Decimal128(est_price)) + stock = encodeKey(stock) + watchlist.stocks[stock] = Decimal128(est_price) watchlist.save() return True def unwatch(user, stock): user_check(user) + stock = stock.replace(".", "\\u002e") watchlist = Watchlist.objects.get(user=user) if stock in watchlist.stocks.keys(): del watchlist.stocks[stock]