polygon_api.py•5.38 kB
import requests
import os
from datetime import datetime, timedelta
from typing import Dict, Optional
from dotenv import load_dotenv
load_dotenv()
class PolygonStockAPI:
"""
A class to interact with Polygon.io API for fetching stock market data.
"""
def __init__(self, api_key: Optional[str] = None):
"""
Initialize the Polygon API client.
Args:
api_key: Polygon.io API key. If not provided, will try to get from environment variable POLYGON_API_KEY
"""
self.api_key = api_key or os.getenv("POLYGON_API_KEY")
if not self.api_key:
raise ValueError("Polygon API key is required. Set POLYGON_API_KEY environment variable or pass api_key parameter.")
self.base_url = "https://api.polygon.io"
def get_last_working_day(self) -> str:
"""
Get the last working day (excluding weekends).
Returns:
Date string in YYYY-MM-DD format
"""
today = datetime.now()
# Go back to find the last working day
days_back = 1
while True:
check_date = today - timedelta(days=days_back)
# Monday = 0, Sunday = 6
if check_date.weekday() < 5: # Monday to Friday
return check_date.strftime("%Y-%m-%d")
days_back += 1
def get_stock_ohlc(self, symbol: str, date: Optional[str] = None) -> Dict:
"""
Fetch Open, High, Low, Close (OHLC) data for a stock symbol.
Args:
symbol: Stock symbol (e.g., 'AAPL', 'MSFT', 'GOOGL')
date: Date in YYYY-MM-DD format. If not provided, uses last working day.
Returns:
Dictionary containing OHLC data and metadata
"""
if not date:
date = self.get_last_working_day()
# Validate symbol format
symbol = symbol.upper().strip()
# Polygon.io API endpoint for daily bars
url = f"{self.base_url}/v2/aggs/ticker/{symbol}/prev"
params = {
"adjusted": "true",
"apikey": self.api_key
}
try:
response = requests.get(url, params=params, timeout=10)
response.raise_for_status()
data = response.json()
if data.get("status") != "OK":
return {
"error": f"API Error: {data.get('status', 'Unknown error')}",
"symbol": symbol,
"date": date
}
results = data.get("results", [])
if not results:
return {
"error": f"No data found for symbol {symbol}",
"symbol": symbol,
"date": date
}
# Extract OHLC data from the first result
result = results[0]
return {
"symbol": symbol,
"date": date,
"open": result.get("o"),
"high": result.get("h"),
"low": result.get("l"),
"close": result.get("c"),
"volume": result.get("v"),
"volume_weighted_average_price": result.get("vw"),
"number_of_transactions": result.get("n"),
"status": "success"
}
except requests.exceptions.RequestException as e:
return {
"error": f"Network error: {str(e)}",
"symbol": symbol,
"date": date
}
except Exception as e:
return {
"error": f"Unexpected error: {str(e)}",
"symbol": symbol,
"date": date
}
def format_stock_data(self, data: Dict) -> str:
"""
Format stock data into a readable string.
Args:
data: Stock data dictionary from get_stock_ohlc
Returns:
Formatted string with stock information
"""
if "error" in data:
return f"Error fetching data for {data.get('symbol', 'Unknown')}: {data['error']}"
return f"""
Stock Symbol: {data['symbol']}
Date: {data['date']}
Open: ${data['open']:.2f}
High: ${data['high']:.2f}
Low: ${data['low']:.2f}
Close: ${data['close']:.2f}
Volume: {data['volume']:,}
VWAP: ${data['volume_weighted_average_price']:.2f}
Transactions: {data['number_of_transactions']:,}
"""
def get_stock_data(symbol: str, date: Optional[str] = None) -> str:
"""
Convenience function to get formatted stock data.
Args:
symbol: Stock symbol (e.g., 'AAPL', 'MSFT', 'GOOGL')
date: Date in YYYY-MM-DD format. If not provided, uses last working day.
Returns:
Formatted string with stock OHLC data
"""
try:
api = PolygonStockAPI()
data = api.get_stock_ohlc(symbol, date)
return api.format_stock_data(data)
except Exception as e:
return f"Error initializing API: {str(e)}"
if __name__ == "__main__":
# Example usage
try:
# Test with Apple stock
result = get_stock_data("AAPL")
print(result)
except Exception as e:
print(f"Error: {e}")