Update core/data.py

core/data.py

@@ -8,19 +8,27 @@ try:
 except ImportError:
     raise ImportError("yfinance must be installed to fetch financial data.")
 
-def load_data(data_src="yahoo", ticker="AAPL", file_upload=None, start="2020-01-01", end="2023-01-01"):
+def load_data(data_src="yahoo", ticker="AAPL", file_upload=None, start="2020-01-01", end="2023-01-01", horizon=1):
+    main_df = None
+    future_df = None
+    
     if data_src == "yahoo":
         try:
             info = yf.Ticker(ticker).info
             if not info:
                 raise ValueError(f"Ticker '{ticker}' not found.")
-            df = yf.download(ticker, start=start, end=end, progress=False)
+            # Fetch data up to end_date + horizon for future actuals
+            extended_end = (pd.to_datetime(end) + pd.Timedelta(days=horizon)).strftime('%Y-%m-%d')
+            df = yf.download(ticker, start=start, end=extended_end, progress=False)
             if df.empty:
                 raise ValueError(f"No data found for ticker '{ticker}' in the specified date range. Please check the symbol and dates.")
+            df = df[['Close']].dropna().rename(columns={'Close': 'value'})
+            df.reset_index(inplace=True)
+            # Split into main_df (up to end_date) and future_df (beyond end_date)
+            main_df = df[df['Date'] <= pd.to_datetime(end)].copy()
+            future_df = df[(df['Date'] > pd.to_datetime(end)) & (df['Date'] <= pd.to_datetime(extended_end))].copy()
         except Exception as e:
             raise ValueError(f"Error fetching data for ticker '{ticker}': {e}")
-        df = df[['Close']].dropna().rename(columns={'Close': 'value'})
-        df.reset_index(inplace=True)
     elif data_src == "csv":
         if file_upload is None:
             raise ValueError("CSV file upload required but not provided.")
@@ -33,10 +41,15 @@ def load_data(data_src="yahoo", ticker="AAPL", file_upload=None, start="2020-01-
                 df = df[['Close']].rename(columns={'Close': 'value'})
             else:
                 raise ValueError("CSV must contain a 'value' or 'Close' column.")
-        df = df[['value']].dropna().reset_index(drop=True)
+        df['Date'] = pd.to_datetime(df.get('Date', df.index))
+        df = df[['Date', 'value']].dropna().reset_index(drop=True)
+        # Split into main_df (up to end_date) and future_df (beyond end_date)
+        main_df = df[df['Date'] <= pd.to_datetime(end)].copy()
+        future_df = df[(df['Date'] > pd.to_datetime(end)) & (df['Date'] <= pd.to_datetime(end) + pd.Timedelta(days=horizon))].copy()
     else:
         raise ValueError("Invalid data source. 'csv' or 'yahoo' expected.")
-    return df
+    
+    return main_df, future_df
 
 
 def preprocess_data(df, column, window_size=30):