MongoDB driver can now handle portfolio functions

- No watchlist functionality yet
This commit is contained in:
socks 2021-03-08 11:59:25 +00:00
parent adf80b0a7b
commit fbfeee8209

56
mongo.py Normal file
View file

@ -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()