From 3a985aa1f563d37107086ddd740769a9e708b832 Mon Sep 17 00:00:00 2001
From: "Bobby (aider)" <bobby@abellana.work>
Date: Mon, 10 Feb 2025 11:09:58 -0800
Subject: [PATCH] feat: Add get_current_prices function with yfinance support

---
 src/utils/data_utils.py | 35 +++++++++++++++++++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/src/utils/data_utils.py b/src/utils/data_utils.py
index 1177de7..c40bba1 100644
--- a/src/utils/data_utils.py
+++ b/src/utils/data_utils.py
@@ -1,5 +1,6 @@
 import os
 import pandas as pd
+import yfinance as yf
 from datetime import datetime, timedelta
 from db.db_connection import create_client
 from screener.user_input import get_interval_choice, get_date_range
@@ -38,6 +39,40 @@ def get_user_input(prompt: str, input_type: type = str, allow_empty: bool = Fals
 def get_float_input(prompt: str) -> Optional[float]:
     return get_user_input(prompt, float)
 
+def get_current_prices(tickers: list) -> dict:
+    """Get current prices for multiple tickers using yfinance"""
+    try:
+        # Create a space-separated string of tickers
+        ticker_str = " ".join(tickers)
+        # Get data for all tickers at once
+        data = yf.download(ticker_str, period="1d", interval="1m", group_by='ticker')
+        
+        prices = {}
+        if len(tickers) == 1:
+            # Handle single ticker case
+            if 'Close' in data.columns:
+                prices[tickers[0]] = data['Close'].iloc[-1]
+            else:
+                print(f"No close price found for {tickers[0]}")
+        else:
+            # Handle multiple tickers
+            for ticker in tickers:
+                try:
+                    if isinstance(data, pd.DataFrame) and (ticker, 'Close') in data.columns:
+                        prices[ticker] = data[ticker]['Close'].iloc[-1]
+                    else:
+                        print(f"No close price found for {ticker}")
+                except Exception as e:
+                    print(f"Error getting price for {ticker}: {e}")
+        
+        return prices
+    except Exception as e:
+        print(f"Error fetching current prices: {e}")
+        print(f"Data structure received: {type(data)}")
+        if isinstance(data, pd.DataFrame):
+            print(f"Available columns: {data.columns}")
+        return {}
+
 def validate_signal_date(signal_date: datetime) -> datetime:
     """
     Validate and adjust signal date if needed