stock-tracker-discord/main.py

194 lines
4.9 KiB
Python

import discord
import table, yfi
import table2 as t2
from discord.ext import commands
import config as cfg
from decimal import Decimal
import typing
if cfg.use_mongodb:
import mongo as database
print("[*] Using MongoDB Database")
else:
import flatfile as database
print("[*] Using FlatFile Database")
intents = discord.Intents.default()
bot = commands.Bot(command_prefix=cfg.command_prefix, intents=intents)
@bot.command()
async def ping(ctx):
await ctx.send("pong!")
@bot.command(aliases=["pft"])
async def portfolio_text(ctx):
user = str(ctx.message.author)
return await ctx.send(
str("```" + table.generate_table(database.get_stocks(user)) + "```")
)
@bot.command(aliases=["pf"])
async def portfolio(ctx):
user = str(ctx.message.author)
portfolio = database.get_stocks(user)
yesterday_portfolio_value = {}
current_portfolio_value = {}
data = []
for stock in portfolio.keys():
delta, old_price = yfi.get_delta(stock, return_old_price=True)
current_price = yfi.get_current_price(stock)
data.append(
[
stock,
Decimal(portfolio[stock]),
round(float(Decimal(portfolio[stock]) * Decimal(current_price)), 2),
str(round(delta, 2)) + "%",
]
)
yesterday_portfolio_value[stock] = Decimal(portfolio[stock]) * Decimal(
old_price
)
current_portfolio_value[stock] = Decimal(portfolio[stock]) * Decimal(
current_price
)
yesterday_portfolio_total_value = sum(
[yesterday_portfolio_value[stock] for stock in yesterday_portfolio_value]
)
current_portfolio_total_value = sum(
[current_portfolio_value[stock] for stock in current_portfolio_value]
)
# TODO: Raises ZeroDivisionError if user has no stocks
total_delta = 100 * (
(current_portfolio_total_value - yesterday_portfolio_total_value)
/ yesterday_portfolio_total_value
)
data.append(
[
"Total",
"",
round(sum([x[2] for x in data]), 2),
str(round(total_delta, 2)) + "%",
]
)
image = t2.generate_table(
headers=["Ticker", "Shares", "Value ($)", "24h"],
data=data,
title=str(user + "'s Portfolio"),
)
await ctx.send(file=discord.File(image, "table.png"))
image.close()
@bot.command()
async def add(ctx, stock, amount):
user = str(ctx.message.author)
stock = stock.upper()
try:
float(amount)
except ValueError:
await ctx.send("amoutn not a number")
return
if database.add_stock(user, stock, amount):
return await ctx.send(
str("```" + table.generate_table(database.get_stocks(user)) + "```")
)
else:
return await ctx.send("Stock **{0}** does not exist!".format(stock))
@bot.command(aliases=["del"])
async def delete(ctx, stock):
stock = stock.upper()
user = str(ctx.message.author)
return await ctx.send(
str("```" + table.generate_table(database.delete_stock(user, stock)) + "```")
)
@bot.command(aliases=["wlt"])
async def watchlist_text(ctx):
user = str(ctx.message.author)
wlist = database.get_watchlist(user)
return await ctx.send(str("```" + table.watchlist_table(wlist) + "```"))
@bot.command(aliases=["w"])
async def watch(ctx, stock, est_price: typing.Optional[int]):
if not est_price:
est_price = 0
user = str(ctx.message.author)
stock = stock.upper()
try:
float(est_price)
except ValueError:
await ctx.send("Estimated price not a number")
return
if database.watch(user, stock, est_price):
return await ctx.send("Updated watchlist")
else:
return await ctx.send("Stock **{0}** does not exist!".format(stock))
@bot.command(aliases=["uw"])
async def unwatch(ctx, *stocks):
user = str(ctx.message.author)
stocks = [stock.upper() for stock in stocks]
for stock in stocks:
database.unwatch(user, stock)
return await ctx.send("Updated watchlist")
@bot.command(aliases=["wl"])
async def watchlist(ctx, user: typing.Optional[discord.Member]):
if not user:
user = str(ctx.message.author)
else:
print(str(user))
user = str(user)
watchlist = database.get_watchlist(user)
data = []
for stock in watchlist.keys():
data.append(
[
stock,
yfi.get_current_price(stock),
Decimal(watchlist[stock]),
str(round(yfi.get_delta(stock), 2)) + "%",
]
)
image = t2.generate_table(
headers=["Ticker", "Value ($)", "Est. Price ($)", "Delta"],
data=data,
title=str(user + "'s Watchlist"),
bg="lightgreen",
border="limegreen",
)
await ctx.send(file=discord.File(image, "table.png"))
image.close()
bot.run(cfg.api_key)