Use Polygon API to get stock price

discord_bot.py

@@ -1,17 +1,21 @@
 import os
 import discord
-
+from polygon import RESTClient
 from dotenv import load_dotenv
 from message_parser import parse_message_content
 import stonk
 
 load_dotenv()
-stonk_bot_token = os.getenv('STONK_BOT_TOKEN')
 
+# initialize Discord API client
+stonk_bot_token = os.getenv('STONK_BOT_TOKEN')
 intents = discord.Intents.default()
 intents.message_content = True
 client = discord.Client(intents=intents)
 
+# initialize Polygon API client
+polygon_api_key = os.getenv('POLYGON_API_KEY')
+polygonClient = RESTClient(api_key=polygon_api_key)
 
 @client.event
 async def on_ready():
@@ -38,4 +42,4 @@ async def on_message(message):
   if symbol == "$help":
     await stonk.handle_help(message)
   else:
-    await stonk.handle_stonk(symbol, message)
+    await stonk.handle_stonk(polygonClient, symbol, message)

from_to.py

@@ -0,0 +1,6 @@
+from datetime import date, timedelta
+
+def from_to():
+    to = date.today()
+    from_ = to - timedelta(days=7)
+    return from_, to

poly.py

@@ -0,0 +1,28 @@
+from from_to import from_to
+
+
+def growth_rate(data):
+  previousClose, lastClose = data[1].close, data[0].close
+  return (lastClose - previousClose) / previousClose
+
+def close_price(data):
+  return data[0].close
+
+def fetch_stock_price_data(stockApiClient, symbol):
+  from_, to = from_to()
+  data = stockApiClient.get_aggs(
+      ticker=symbol,
+      multiplier=1,
+      timespan="day",
+      from_=from_,
+      to=to,
+      adjusted=True,
+      sort="desc",
+      limit=50,
+  )
+  info = stockApiClient.get_ticker_details(symbol)
+
+  price = close_price(data)
+  growth = growth_rate(data)
+  name = info.name
+  return price, growth, name

requirements.txt

@@ -19,6 +19,7 @@ multidict==6.0.4
 multitasking==0.0.11
 numpy==1.24.1
 pandas==1.5.2
+polygon-api-client==1.7.1
 pycparser==2.21
 python-dateutil==2.8.2
 python-dotenv==0.21.0
@@ -28,5 +29,6 @@ six==1.16.0
 soupsieve==2.3.2.post1
 urllib3==1.26.14
 webencodings==0.5.1
+websockets==10.4
 yarl==1.8.2
 yfinance==0.2.9

stonk.py

@@ -1,44 +1,53 @@
-import yfinance as yf
-
+from from_to import from_to
 from number_formatter import format_growth_rate, format_price
+import poly
+import yahoo
 
+HELP_CONTENT = """
+  $stonk [SYMBOL]
+  
+  Examples:
+  - $stonk $help: show this help message
+  - $stonk AAPL: show the latest stock price of AAPL
+"""
 
-def growth_rate(data):
-  previouseClose, lastClose = data['Close'][-2], data['Close'][-1]
-  return (lastClose - previouseClose) / previouseClose
+def use_polygon(client, symbol):
+  try:
+    price, growth, name = poly.fetch_stock_price_data(client, symbol)
+  except Exception as error:
+    print("Failed to fetch data from Polygon {}".format(error))
+    return None, False
+  return (price, growth, name), True
 
-def close_price(data):
-  return data['Close'][-1]
+def use_yahoo(symbol):
+  try:
+    price, growth, name = yahoo.fetch_stock_price_data(symbol)
+  except Exception as error:
+    print("Failed to fetch data from Yahoo! Finance {}".format(error))
+    return None, False
+  return (price, growth, name), True
 
-def fetch_price_data(symbol):
-  ticker = yf.Ticker(symbol)
-  data = ticker.history(period='7d')
-  info = ticker.info
-  return data, info
+async def reply(message, results, symbol):
+  price, growth, name = results
+  await message.reply("{} ({}): {} {}".format(symbol, name, format_price(price), format_growth_rate(growth)))
 
-async def handle_stonk(symbol, message):
+async def handle_stonk(stockApiClient, symbol, message):
   if symbol is None:
     await message.reply("Ticker symbol cannot be empty")
     return
 
-  try:
-    data, info = fetch_price_data(symbol)
-    price = close_price(data)
-    growth = growth_rate(data)
-    name = info["shortName"]
-    await message.reply("{} ({}): {} {}".format(symbol, name, format_price(price), format_growth_rate(growth)))
-  except Exception as error:
-    print(error)
-    await message.reply("Cannot get data for {}".format(symbol))
+  results, success = use_polygon(stockApiClient, symbol)
+  if success:
+    await reply(message, results, symbol)
+    return
+
+  results, success = use_yahoo(symbol)
+  if success:
+    await reply(message, results, symbol)
     return
+
+  await message.reply("Cannot get data for {}".format(symbol))
 
 
 async def handle_help(message):
-  content = """
-  $stonk [SYMBOL]
-  
-  Examples:
-  - $stonk $help: show this help message
-  - $stonk AAPL: show the latest stock price of AAPL
-"""
-  await message.reply(content)
+  await message.reply(HELP_CONTENT)

yahoo.py

@@ -0,0 +1,18 @@
+import yfinance as yf
+
+def growth_rate(data):
+  previouseClose, lastClose = data['Close'][-2], data['Close'][-1]
+  return (lastClose - previouseClose) / previouseClose
+
+def close_price(data):
+  return data['Close'][-1]
+
+def fetch_stock_price_data(symbol):
+  ticker = yf.Ticker(symbol)
+  data = ticker.history(period='7d')
+  info = ticker.info
+
+  price = close_price(data)
+  growth = growth_rate(data)
+  name = info.name
+  return price, growth, name