Source code for odd_kernel.datasets.marketdata

import pandas as pd
import yfinance as yf
from enum import Enum
from typing import List, Dict, Union
from curl_cffi import requests

NANOSECONDS_PER_SECOND = 1_000_000_000

[docs] class DataType(Enum): """Available financial data fields.""" OPEN = "Open" HIGH = "High" LOW = "Low" CLOSE = "Close" ADJ_CLOSE = "Adj Close" VOLUME = "Volume"
[docs] @classmethod def parse(cls, value: str): """ Parses a string (case-insensitive) into a DataType. Raises ValueError if the string does not match any field. """ normalized = value.strip().lower() for field in cls: if field.value.lower() == normalized: return field raise ValueError(f"Unknown data field: {value}")
[docs] class MarketDataTimeUnit(Enum): """Supported time units for interpolation resolution.""" DAY = "D" MONTH = "M" YEAR = "Y"
[docs] class IndexType(Enum): """Supported index types for the output data.""" STRING = "string" DATETIME = "datetime" EPOCH = "epoch" INDEX = "index"
[docs] class MarketDataPeriod: """Represents a time resolution with unit and magnitude.""" def __init__(self, value: int, unit: MarketDataTimeUnit): if value <= 0: raise ValueError("Period value must be positive.") self.value = value self.unit = unit
[docs] def to_pandas_freq(self) -> str: """Converts the period to a pandas-compatible frequency string.""" return f"{self.value}{self.unit.value}"
def __str__(self): return self.to_pandas_freq()
[docs] class MarketDataProvider: """ Module to fetch and manage financial time series from Yahoo Finance. Provides summary, raw data, and interpolated data methods. """
[docs] def get_available_tickers(self, max_count = None): """ Returns all available tickers in the yahoo finance API """ offset = 0 batch_size = 100 all_tickers = [] while True: url = f"https://query1.finance.yahoo.com/v1/finance/screener/predefined/saved?count={batch_size}&formatted=false&scrIds=MOST_EXPENSIVE_STOCKS&sortField=dayvolume&sortType=DESC&start={offset}&useRecordsResponse=true&fields=ticker" response = requests.get(url, impersonate="chrome") response.raise_for_status() response_body = response.json() # Parse tickers tickers_info = response_body['finance']['result'][0]['records'] all_tickers.extend(ticker_info['ticker'] for ticker_info in tickers_info) if len(tickers_info) < batch_size or max_count is not None and max_count <= len(all_tickers): break offset += batch_size return all_tickers[:max_count]
def __init__(self, max_cache_size=100, show_download_progress=False): """ Initializes the MarketDataProvider. Args: max_cache_size (int, optional): The maximum number of tickers to cache. Defaults to 100. show_download_progress (bool, optional): Whether to show the download progress. Defaults to False. """ self.dates_map = {} self.data_cache = {} self.metatadata_cache = {} self.max_cache_size = max_cache_size self.cache_keys = [] self.show_download_progress = show_download_progress def _download(self, tickers: Union[str, List[str]]) -> tuple[Dict, Dict]: """ Downloads and caches Yahoo Finance data for one or multiple tickers. Args: tickers (Union[str, List[str]]): A single ticker or a list of tickers. Returns: tuple[Dict, Dict]: A tuple containing two dictionaries: one for the data and one for the metadata. """ if isinstance(tickers, str): tickers = [tickers] missing = [] data_by_ticker = {} meta_data_by_ticker = {} for ticker in tickers: if ticker in self.data_cache: data_by_ticker[ticker] = self.data_cache[ticker] meta_data_by_ticker[ticker] = self.metatadata_cache[ticker] else: missing.append(ticker) if missing: data = yf.download( missing, start="1800-12-31", end=pd.Timestamp.today().strftime("%Y-%m-%d"), progress=self.show_download_progress, auto_adjust=False, multi_level_index=False ) if data.empty: raise ValueError(f"No data retrieved for \'{missing}\'.") if isinstance(data.columns, pd.MultiIndex): data_types = set(DataType.parse(data_type) for data_type, _ in data.columns) for ticker in missing: data_by_ticker[ticker] = {} meta_data_by_ticker[ticker] = yf.Ticker(ticker).info for data_type in data_types: data_by_ticker[ticker][data_type] = data[(data_type.value, ticker)].dropna() self._add_to_cache(ticker, data_by_ticker[ticker], meta_data_by_ticker[ticker]) else: ticker = missing[0] data_by_ticker[ticker] = {} meta_data_by_ticker[ticker] = yf.Ticker(ticker).info data_types = set(DataType.parse(data_type) for data_type in data.columns) for data_type in data_types: data_by_ticker[ticker][data_type] = data[data_type.value].dropna() self._add_to_cache(ticker, data_by_ticker[ticker], meta_data_by_ticker[ticker]) return data_by_ticker, meta_data_by_ticker def _add_to_cache(self, key, data: dict, metadata: dict): """Helper to manage cache insertion.""" if self.max_cache_size <= len(self.cache_keys): key_to_pop = self.cache_keys[0] self.data_cache.pop(key_to_pop) self.metatadata_cache.pop(key_to_pop) self.cache_keys = self.cache_keys[1:] self.data_cache[key] = data self.metatadata_cache[key] = metadata self.cache_keys.append(key)
[docs] def get_ticker_metadata(self, tickers : Union[str, List[str]]) -> Dict: """ Returns the metadata for a set of tickers. """ if type(tickers) is str: return yf.Ticker(tickers).info else: return {ticker:yf.Ticker(ticker).info for ticker in tickers}
[docs] def get_summary(self, tickers: Union[str, List[str]]) -> pd.DataFrame: """ Returns a summary for each ticker: - name - min/max date - average resolution - available fields """ summaries = [] data_by_ticker, _ = self._download(tickers) for ticker, data in data_by_ticker.items(): df = data[next(iter(data))] deltas = df.index.to_series().diff().dropna() resolution = deltas.mode()[0] if not deltas.empty else pd.Timedelta("NaT") summaries.append( { "name": ticker, "date_min": df.index.min().date(), "date_max": df.index.max().date(), "resolution": str(resolution), "fields": list(data.keys()), } ) return pd.DataFrame(summaries)
[docs] def get_raw(self, tickers: Union[str, List[str]], field: DataType) -> tuple[Dict, Dict]: """Returns raw data for a given ticker and field.""" data_by_ticker, metadata_by_ticker = self._download(tickers) return {ticker: data_by_ticker[ticker][field] for ticker in data_by_ticker.keys()}, metadata_by_ticker
[docs] def get_interpolated( self, tickers: Union[str, List[str]], start: str, end: str, field: DataType, resolution: MarketDataPeriod = None, index_type: IndexType = IndexType.DATETIME, extrapolate_left=False, ) -> tuple[Dict, Dict]: """ Returns interpolated data for a given ticker, field, and resolution. Performs flat extrapolation to the left and right of the original data range if required. Args: tickers (Union[str, List[str]]): Ticker symbols of the financial instrument. start (str): Start date (inclusive) in 'YYYY-MM-DD' format. end (str): End date (inclusive) in 'YYYY-MM-DD' format. field (DataType): The financial data field to retrieve (e.g., CLOSE, OPEN). resolution (MarketDataPeriod, optional): The target temporal resolution. If None no interpolation is performed. Defaults to None. index_type (IndexType, optional): Defines the format of the time index in the output. Defaults to IndexType.DATETIME. extrapolate_left (bool, optional): Whether to perform flat extrapolation to the left or not. Defaults to False. Returns: tuple[Dict, Dict]: A tuple containing two dictionaries: one for the interpolated data and one for the metadata. """ data_by_ticker, metadata_by_ticker = self.get_raw(tickers, field) if resolution is not None: freq = resolution.to_pandas_freq() full_index = pd.date_range(start=start, end=end, freq=freq) interpolated_data = {} for ticker in data_by_ticker: interpolated_data[ticker] = data_by_ticker[ticker].reindex(full_index).interpolate(method="time").ffill() if extrapolate_left: interpolated_data[ticker] = interpolated_data[ticker].bfill() else: interpolated_data = {} date_max = pd.to_datetime(end) date_min = pd.to_datetime(start) for ticker, data in data_by_ticker.items(): mask = (date_min <= data.index) & (data.index <= date_max) interpolated_data[ticker] = data.loc[mask] for ticker, interpolated in interpolated_data.items(): if index_type == IndexType.STRING: interpolated_data[ticker].index = interpolated.index.strftime("%Y-%m-%d") elif index_type == IndexType.EPOCH: interpolated_data[ticker].index = ( interpolated.index.astype("int64") / NANOSECONDS_PER_SECOND ) elif index_type == IndexType.INDEX: interpolated_data[ticker].index = range(len(interpolated)) interpolated_data[ticker].index.name = "date" return interpolated_data, metadata_by_ticker
[docs] def get_dataset( self, tickers: List[str], field: DataType, start: str, end: str, resolution: MarketDataPeriod, index_type: IndexType = IndexType.DATETIME, ) -> Dict[str, pd.DataFrame]: """Returns interpolated data for the given tickers, field and resolution in a common time grid""" return { ticker: self.get_interpolated(tickers, start, end, field, resolution, index_type) for ticker in tickers }
# To update this list, run the following script in the browser console on https://en.wikipedia.org/wiki/List_of_S%26P_500_companies # (() => { # const rows = document.querySelectorAll("#constituents tbody tr"); # const tickers = Array.from(rows) # .map(row => row.querySelector("td a")?.textContent.trim()) # .filter(Boolean); # console.log("[\n '" + tickers.join("', '") + "'\n]"); # })(); AVAILABLE_TICKERS = ['MMM', 'AOS', 'ABT', 'ABBV', 'ACN', 'ADBE', 'AMD', 'AES', 'AFL', 'A', 'APD', 'ABNB', 'AKAM', 'ALB', 'ARE', 'ALGN', 'ALLE', 'LNT', 'ALL', 'GOOGL', 'GOOG', 'MO', 'AMZN', 'AMCR', 'AEE', 'AEP', 'AXP', 'AIG', 'AMT', 'AWK', 'AMP', 'AME', 'AMGN', 'APH', 'ADI', 'AON', 'APA', 'APO', 'AAPL', 'AMAT', 'APP', 'APTV', 'ACGL', 'ADM', 'ANET', 'AJG', 'AIZ', 'T', 'ATO', 'ADSK', 'ADP', 'AZO', 'AVB', 'AVY', 'AXON', 'BKR', 'BALL', 'BAC', 'BAX', 'BDX', 'BBY', 'TECH', 'BIIB', 'BLK', 'BX', 'XYZ', 'BK', 'BA', 'BKNG', 'BSX', 'BMY', 'AVGO', 'BR', 'BRO', 'BLDR', 'BG', 'BXP', 'CHRW', 'CDNS', 'CPT', 'CPB', 'COF', 'CAH', 'KMX', 'CCL', 'CARR', 'CAT', 'CBOE', 'CBRE', 'CDW', 'COR', 'CNC', 'CNP', 'CF', 'CRL', 'SCHW', 'CHTR', 'CVX', 'CMG', 'CB', 'CHD', 'CI', 'CINF', 'CTAS', 'CSCO', 'C', 'CFG', 'CLX', 'CME', 'CMS', 'KO', 'CTSH', 'COIN', 'CL', 'CMCSA', 'CAG', 'COP', 'ED', 'STZ', 'CEG', 'COO', 'CPRT', 'GLW', 'CPAY', 'CTVA', 'CSGP', 'COST', 'CTRA', 'CRWD', 'CCI', 'CSX', 'CMI', 'CVS', 'DHR', 'DRI', 'DDOG', 'DVA', 'DAY', 'DECK', 'DE', 'DELL', 'DAL', 'DVN', 'DXCM', 'FANG', 'DLR', 'DG', 'DLTR', 'D', 'DPZ', 'DASH', 'DOV', 'DOW', 'DHI', 'DTE', 'DUK', 'DD', 'EMN', 'ETN', 'EBAY', 'ECL', 'EIX', 'EW', 'EA', 'ELV', 'EME', 'EMR', 'ETR', 'EOG', 'EPAM', 'EQT', 'EFX', 'EQIX', 'EQR', 'ERIE', 'ESS', 'EL', 'EG', 'EVRG', 'ES', 'EXC', 'EXE', 'EXPE', 'EXPD', 'EXR', 'XOM', 'FFIV', 'FDS', 'FICO', 'FAST', 'FRT', 'FDX', 'FIS', 'FITB', 'FSLR', 'FE', 'FI', 'F', 'FTNT', 'FTV', 'FOXA', 'FOX', 'BEN', 'FCX', 'GRMN', 'IT', 'GE', 'GEHC', 'GEV', 'GEN', 'GNRC', 'GD', 'GIS', 'GM', 'GPC', 'GILD', 'GPN', 'GL', 'GDDY', 'GS', 'HAL', 'HIG', 'HAS', 'HCA', 'DOC', 'HSIC', 'HSY', 'HPE', 'HLT', 'HOLX', 'HD', 'HON', 'HRL', 'HST', 'HWM', 'HPQ', 'HUBB', 'HUM', 'HBAN', 'HII', 'IBM', 'IEX', 'IDXX', 'ITW', 'INCY', 'IR', 'PODD', 'INTC', 'IBKR', 'ICE', 'IFF', 'IP', 'IPG', 'INTU', 'ISRG', 'IVZ', 'INVH', 'IQV', 'IRM', 'JBHT', 'JBL', 'JKHY', 'J', 'JNJ', 'JCI', 'JPM', 'K', 'KVUE', 'KDP', 'KEY', 'KEYS', 'KMB', 'KIM', 'KMI', 'KKR', 'KLAC', 'KHC', 'KR', 'LHX', 'LH', 'LRCX', 'LW', 'LVS', 'LDOS', 'LEN', 'LII', 'LLY', 'LIN', 'LYV', 'LKQ', 'LMT', 'L', 'LOW', 'LULU', 'LYB', 'MTB', 'MPC', 'MAR', 'MMC', 'MLM', 'MAS', 'MA', 'MTCH', 'MKC', 'MCD', 'MCK', 'MDT', 'MRK', 'META', 'MET', 'MTD', 'MGM', 'MCHP', 'MU', 'MSFT', 'MAA', 'MRNA', 'MHK', 'MOH', 'TAP', 'MDLZ', 'MPWR', 'MNST', 'MCO', 'MS', 'MOS', 'MSI', 'MSCI', 'NDAQ', 'NTAP', 'NFLX', 'NEM', 'NWSA', 'NWS', 'NEE', 'NKE', 'NI', 'NDSN', 'NSC', 'NTRS', 'NOC', 'NCLH', 'NRG', 'NUE', 'NVDA', 'NVR', 'NXPI', 'ORLY', 'OXY', 'ODFL', 'OMC', 'ON', 'OKE', 'ORCL', 'OTIS', 'PCAR', 'PKG', 'PLTR', 'PANW', 'PSKY', 'PH', 'PAYX', 'PAYC', 'PYPL', 'PNR', 'PEP', 'PFE', 'PCG', 'PM', 'PSX', 'PNW', 'PNC', 'POOL', 'PPG', 'PPL', 'PFG', 'PG', 'PGR', 'PLD', 'PRU', 'PEG', 'PTC', 'PSA', 'PHM', 'PWR', 'QCOM', 'DGX', 'RL', 'RJF', 'RTX', 'O', 'REG', 'REGN', 'RF', 'RSG', 'RMD', 'RVTY', 'HOOD', 'ROK', 'ROL', 'ROP', 'ROST', 'RCL', 'SPGI', 'CRM', 'SBAC', 'SLB', 'STX', 'SRE', 'NOW', 'SHW', 'SPG', 'SWKS', 'SJM', 'SW', 'SNA', 'SOLV', 'SO', 'LUV', 'SWK', 'SBUX', 'STT', 'STLD', 'STE', 'SYK', 'SMCI', 'SYF', 'SNPS', 'SYY', 'TMUS', 'TROW', 'TTWO', 'TPR', 'TRGP', 'TGT', 'TEL', 'TDY', 'TER', 'TSLA', 'TXN', 'TPL', 'TXT', 'TMO', 'TJX', 'TKO', 'TTD', 'TSCO', 'TT', 'TDG', 'TRV', 'TRMB', 'TFC', 'TYL', 'TSN', 'USB', 'UBER', 'UDR', 'ULTA', 'UNP', 'UAL', 'UPS', 'URI', 'UNH', 'UHS', 'VLO', 'VTR', 'VLTO', 'VRSN', 'VRSK', 'VZ', 'VRTX', 'VTRS', 'VICI', 'V', 'VST', 'VMC', 'WRB', 'GWW', 'WAB', 'WMT', 'DIS', 'WBD', 'WM', 'WAT', 'WEC', 'WFC', 'WELL', 'WST', 'WDC', 'WY', 'WSM', 'WMB', 'WTW', 'WDAY', 'WYNN', 'XEL', 'XYL', 'YUM', 'ZBRA', 'ZBH', 'ZTS']