diff --git a/app/cron_similar_stocks.py b/app/cron_similar_stocks.py index c5e838d..631a5a7 100755 --- a/app/cron_similar_stocks.py +++ b/app/cron_similar_stocks.py @@ -1,66 +1,89 @@ -import ujson +import orjson import asyncio import aiohttp import sqlite3 from datetime import datetime -from rating import rating_model import pandas as pd from tqdm import tqdm -async def save_similar_stocks(symbol, data): - with open(f"json/similar-stocks/{symbol}.json", 'w') as file: - ujson.dump(data, file) +async def save_json(symbol, data): + with open(f"json/similar-stocks/{symbol}.json", 'wb') as file: + file.write(orjson.dumps(data)) +# Load stock screener data +with open(f"json/stock-screener/data.json", 'rb') as file: + stock_screener_data = orjson.loads(file.read()) +stock_screener_data_dict = {item['symbol']: item for item in stock_screener_data} query_template = """ SELECT - quote, stock_peers + stock_peers FROM stocks WHERE symbol = ? """ +async def get_data(symbol): + """Extract specified columns data for a given symbol.""" + columns = ['dividendYield', 'employees', 'marketCap'] + + if symbol in stock_screener_data_dict: + result = {} + for column in columns: + result[column] = stock_screener_data_dict[symbol].get(column, None) + return result + return {} + async def run(): cursor = con.cursor() cursor.execute("PRAGMA journal_mode = wal") - cursor.execute("SELECT DISTINCT symbol FROM stocks") - stocks_symbols = [row[0] for row in cursor.fetchall()] - for ticker in tqdm(stocks_symbols): - filtered_df = [] + cursor.execute("SELECT DISTINCT symbol FROM stocks WHERE symbol NOT LIKE '%.%'") + total_symbols = [row[0] for row in cursor.fetchall()] + #total_symbols = ['NVDA'] # For testing purposes + + for ticker in tqdm(total_symbols): + # Get peers for the current ticker df = pd.read_sql_query(query_template, con, params=(ticker,)) try: - df = ujson.loads(df['stock_peers'].iloc[0]) + # Get the list of peer stocks + peers = orjson.loads(df['stock_peers'].iloc[0]) + + # Create a list to store peer data + peer_data_list = [] + + # Process each peer + for peer_symbol in peers: + # Get additional data for this peer + data = await get_data(peer_symbol) + + # Combine symbol with additional data + peer_info = { + 'symbol': peer_symbol, + **data + } + peer_data_list.append(peer_info) + + # Sort by marketCap if available + sorted_peers = sorted( + peer_data_list, + key=lambda x: x.get('marketCap', 0) or 0, + reverse=True + ) + + # Save the results + if sorted_peers: + await save_json(ticker, sorted_peers) + except: - df = [] - if len(df) > 0: - df = [stock for stock in df if stock in stocks_symbols] - for symbol in df: - try: - df = pd.read_sql_query(query_template, con, params=(symbol,)) - df_dict = df.to_dict() - quote_dict = eval(df_dict['quote'][0])[0] - filtered_df.append(quote_dict) # Add the modified result to the combined list - except: - pass + pass - filtered_df = [ - { - "symbol": entry["symbol"], - "name": entry["name"], - "marketCap": entry["marketCap"], - "avgVolume": entry["avgVolume"] - } - for entry in filtered_df - ] - - sorted_df = sorted(filtered_df, key=lambda x: x['marketCap'], reverse=True) - - await save_similar_stocks(ticker, sorted_df) - -try: - con = sqlite3.connect('stocks.db') - asyncio.run(run()) - con.close() -except Exception as e: - print(e) \ No newline at end of file +if __name__ == "__main__": + try: + con = sqlite3.connect('stocks.db') + asyncio.run(run()) + except Exception as e: + print(f"Error: {e}") + finally: + if 'con' in locals(): + con.close() \ No newline at end of file