stock-tracker-discord/database.py

50 lines
1.3 KiB
Python

import json
import asyncio
import yfinance as yf
from decimal import Decimal
from config import config
db = config['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)
print("USER" + str(user))
if user in list(data.keys()):
user_data = data[user]
return user_data['portfolio']
else:
print(data)
data[user] = {"portfolio": {}}
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()))
print(data)
print(type(data))
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)