update cron job for similar stocks

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)
+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()