116 lines
2.6 KiB
Python
116 lines
2.6 KiB
Python
import json
|
|
import asyncio
|
|
import yfinance as yf
|
|
from decimal import Decimal
|
|
import config as cfg
|
|
|
|
db = cfg.database
|
|
|
|
|
|
def stock_exists(ticker):
|
|
try:
|
|
yf.Ticker(ticker).info
|
|
return True
|
|
except KeyError:
|
|
return False
|
|
|
|
|
|
def write_file(data):
|
|
with open(db, "w") as f:
|
|
f.write(data)
|
|
|
|
|
|
def get_stocks(user):
|
|
with open(db) as f:
|
|
data = json.loads(f.read())
|
|
|
|
data = dict(data)
|
|
if user in list(data.keys()):
|
|
user_data = data[user]
|
|
return user_data["portfolio"]
|
|
else:
|
|
data[user] = {"portfolio": {}, "watchlist": {}}
|
|
write_file(json.dumps(data))
|
|
return get_stocks(user)
|
|
|
|
|
|
def add_stock(user, stock, amount):
|
|
|
|
with open(db) as f:
|
|
data = dict(json.loads(f.read()))
|
|
if user in data.keys():
|
|
if stock in data[user]["portfolio"].keys():
|
|
data[user]["portfolio"][stock] = str(
|
|
Decimal(amount) + Decimal(data[user]["portfolio"][stock])
|
|
)
|
|
elif stock_exists(stock):
|
|
data[user]["portfolio"][stock] = str(Decimal(amount))
|
|
else:
|
|
return False
|
|
write_file(json.dumps(data))
|
|
return True
|
|
else:
|
|
get_stocks(user)
|
|
return add_stock(user, stock, amount)
|
|
|
|
|
|
def delete_stock(user, stock):
|
|
with open(db) as f:
|
|
data = json.loads(f.read())
|
|
|
|
if user in data.keys():
|
|
# user exists
|
|
portfolio = data[user]["portfolio"]
|
|
if stock in portfolio.keys():
|
|
# stock exists
|
|
del portfolio[stock]
|
|
data[user]["portfolio"] = portfolio
|
|
write_file(json.dumps(data))
|
|
return get_stocks(user)
|
|
|
|
|
|
def get_watchlist(user):
|
|
with open(db) as f:
|
|
data = json.loads(f.read())
|
|
|
|
if user not in data.keys():
|
|
get_stocks(user)
|
|
return get_watchlist(user)
|
|
|
|
return data[user]["watchlist"]
|
|
|
|
|
|
def watch(user, stock, est_price=0):
|
|
with open(db) as f:
|
|
data = json.loads(f.read())
|
|
|
|
if user not in data.keys():
|
|
get_stocks(user)
|
|
return add_to_watchlist(user, list_name, stock)
|
|
|
|
if not stock_exists(stock):
|
|
return False
|
|
|
|
watchlist = data[user]["watchlist"]
|
|
watchlist[stock] = str(Decimal(est_price))
|
|
data[user]["watchlist"] = watchlist
|
|
write_file(json.dumps(data))
|
|
|
|
return True
|
|
|
|
|
|
def unwatch(user, stock):
|
|
with open(db) as f:
|
|
data = json.loads(f.read())
|
|
|
|
if user not in data.keys():
|
|
get_stocks(user)
|
|
return
|
|
|
|
if not stock_exists:
|
|
return
|
|
|
|
watchlist = data[user]["watchlist"]
|
|
del watchlist[stock]
|
|
data[user]["watchlist"] = watchlist
|
|
write_file(json.dumps(data))
|