diff --git a/mongo.py b/mongo.py new file mode 100644 index 0000000..5d549f7 --- /dev/null +++ b/mongo.py @@ -0,0 +1,56 @@ +from mongoengine import * +from config import config +from bson import Decimal128 +from decimal import Decimal +import yfi + +connect(config["database"]) + + +class Portfolio(Document): + user = StringField(required=True) + stocks = DictField(default={}) + + +class Watchlist(Document): + user = StringField(max_length=200, required=True) + stocks = DictField(required=True, default={}) + + +def user_check(user): + # Will create empty portfolio for new users + if not Portfolio.objects(user=user): + portfolio = Portfolio(user=user) + portfolio.save() + + + +# Portfolio Stuff +def get_stocks(user): + user_check(user) + return Portfolio.objects.get(user=user).stocks + +def add_stock(user, stock, amount): + user_check(user) + amount = Decimal(amount) + portfolio = Portfolio.objects.get(user=user) + + if not ( se := yfi.stock_exists(stock)): + # Stock ticker does not exist on yahoo finance + return False + + if stock in portfolio.stocks.keys(): + portfolio.stocks[stock] = Decimal128( + amount + portfolio.stocks[stock].to_decimal() + ) + else: + portfolio.stocks[stock] = Decimal128(amount) + portfolio.save() + return True + +def delete_stock(user, stock): + user_check(user) + portfolio = Portfolio.objects.get(user=user) + if stock in portfolio.stocks.keys(): + del portfolio.stocks[stock] + portfolio.save()