From a0aad6bc9738dfdd38c530322c909f8651faf080 Mon Sep 17 00:00:00 2001
From: socks <me@sockspls.com>
Date: Mon, 8 Mar 2021 12:46:43 +0000
Subject: [PATCH] Added watchlist functionality

---
 mongo.py | 44 ++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 40 insertions(+), 4 deletions(-)

diff --git a/mongo.py b/mongo.py
index 5d549f7..d677c8c 100644
--- a/mongo.py
+++ b/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