Added watchlist functionality

mongo.py

@@ -14,28 +14,33 @@ class Portfolio(Document):
 
 class Watchlist(Document):
     user = StringField(max_length=200, required=True)
-    stocks = DictField(required=True, default={})
+    stocks = DictField(default={})
 
 
 def user_check(user):
-    # Will create empty portfolio for new users
+    # Will create empty portfolio and watchlist for new users
     if not Portfolio.objects(user=user):
         portfolio = Portfolio(user=user)
         portfolio.save()
-
+    if not Watchlist.objects(user=user):
+        watchlist = Watchlist(user=user)
+        watchlist.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)):
+    if not (yfi.stock_exists(stock)):
         # Stock ticker does not exist on yahoo finance
         return False
 
@@ -48,9 +53,40 @@ def add_stock(user, stock, 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()
+
+
+# Watchlist Stuff
+
+
+def get_watchlist(user):
+    user_check(user)
+    return Watchlist.objects.get(user=user).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))
+    watchlist.save()
+    return True
+
+
+def unwatch(user, stock):
+    user_check(user)
+    watchlist = Watchlist.objects.get(user=user)
+    if stock in watchlist.stocks.keys():
+        del watchlist.stocks[stock]
+        watchlist.save()
+    return True