stock-tracker-discord/main.py

import discord
import table, yfi
import table2 as t2
from discord.ext import commands
import config as cfg
from decimal import Decimal
import typing
import inspect

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=["stock", "t", "s"])
async def ticker(ctx, stock):
    user = str(ctx.message.author)
    stock = stock.upper()

    data = yfi.stock_info(stock)
    await ctx.send(
        inspect.cleandoc(
            """**{ticker}**

        Bid: {bid}
        Ask: {ask}
        """.format(
                ticker=stock, bid=data["bid"], ask=data["ask"]
            )
        )
    )

    image = t2.generate_chart(stock, yfi.stock_history(stock))

    await ctx.send(file=discord.File(image, "chart.png"))
    image.close()


@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 ($)", "24h"],
        data=data,
        title=str(user + "'s Watchlist"),
        bg="lightgreen",
        border="limegreen",
    )

    await ctx.send(file=discord.File(image, "table.png"))
    image.close()


@bot.command(aliases=["opt"])
async def options(ctx, stock, money, exp_price):
    user = str(ctx.message.author)
    stock = stock.upper()
    try:
        money = float(money)
        exp_price = float(exp_price)
    except ValueError:
        return await ctx.send("amoutn not a number")

    best_price, prices = yfi.options_chain(stock, money, exp_price)

    output = inspect.cleandoc("```")
    for price in prices:
        if price[0] == best_price[0]:
            output = (
                output
                + "\n"
                + "* Buy strike ${} to make ${:.2f}".format(price[0], price[1])
            )
        elif price[1] < -5000:
            pass
        else:
            output = (
                output
                + "\n"
                + "  Buy strike ${} to make ${:.2f}".format(price[0], price[1])
            )
    output = output + "```"
    print(len(output))
    return await ctx.send(output)


bot.run(cfg.api_key)