From a8dd33c3d9196eeee08b0a6934f7641966f7502f Mon Sep 17 00:00:00 2001
From: "Bobby (aider)" <bobby@abellana.work>
Date: Sat, 8 Feb 2025 12:19:34 -0800
Subject: [PATCH] feat: add date range input functionality to stock scanners

---
 src/screener/t_atr_ema.py    |  5 ++--
 src/screener/t_atr_ema_v2.py |  5 ++--
 src/screener/user_input.py   | 45 ++++++++++++++++++++++++++++++++++++
 3 files changed, 49 insertions(+), 6 deletions(-)

diff --git a/src/screener/t_atr_ema.py b/src/screener/t_atr_ema.py
index d3a59e1..3b6e495 100644
--- a/src/screener/t_atr_ema.py
+++ b/src/screener/t_atr_ema.py
@@ -1,4 +1,4 @@
-from screener.user_input import get_interval_choice
+from screener.user_input import get_interval_choice, get_date_range
 import os
 from datetime import datetime, timedelta
 import pandas as pd
@@ -47,8 +47,7 @@ def run_atr_ema_scanner(min_price: float, max_price: float, min_volume: int, por
     # Get time interval
     interval = get_interval_choice()
     
-    end_date = datetime.now()
-    start_date = end_date - timedelta(days=1)  # Get last trading day
+    start_date, end_date = get_date_range()
     start_ts = int(start_date.timestamp() * 1000000000)
     end_ts = int(end_date.timestamp() * 1000000000)
     
diff --git a/src/screener/t_atr_ema_v2.py b/src/screener/t_atr_ema_v2.py
index 0ee2a50..4b20adc 100644
--- a/src/screener/t_atr_ema_v2.py
+++ b/src/screener/t_atr_ema_v2.py
@@ -4,7 +4,7 @@ import os
 from db.db_connection import create_client
 from trading.position_calculator import PositionCalculator
 from utils.data_utils import get_stock_data
-from screener.user_input import get_interval_choice
+from screener.user_input import get_interval_choice, get_date_range
 from indicators.three_atr_ema import ThreeATREMAIndicator
 
 def check_entry_signal(df: pd.DataFrame) -> bool:
@@ -61,8 +61,7 @@ def run_atr_ema_scanner_v2(min_price: float, max_price: float, min_volume: int,
     
     interval = get_interval_choice()
     
-    end_date = datetime.now()
-    start_date = end_date - timedelta(days=1)
+    start_date, end_date = get_date_range()
     start_ts = int(start_date.timestamp() * 1000000000)
     end_ts = int(end_date.timestamp() * 1000000000)
     
diff --git a/src/screener/user_input.py b/src/screener/user_input.py
index ac1c801..38ff318 100644
--- a/src/screener/user_input.py
+++ b/src/screener/user_input.py
@@ -1,5 +1,50 @@
+from datetime import datetime, timedelta
 from screener.screeners import SCREENERS  # Import SCREENERS dictionary
 
+def get_date_input(prompt: str, default_date: datetime = None) -> datetime:
+    """
+    Get a date input from the user with validation
+    
+    Args:
+        prompt (str): Message to display to user
+        default_date (datetime, optional): Default date if user presses enter
+        
+    Returns:
+        datetime: The validated date
+    """
+    while True:
+        date_str = input(prompt)
+        if not date_str and default_date:
+            return default_date
+            
+        try:
+            return datetime.strptime(date_str, "%Y-%m-%d")
+        except ValueError:
+            print("Invalid date format. Please use YYYY-MM-DD")
+
+def get_date_range() -> tuple:
+    """
+    Get start and end dates from user
+    
+    Returns:
+        tuple: (start_date, end_date) as datetime objects
+    """
+    print("\nEnter date range (format: YYYY-MM-DD)")
+    print("Press Enter for end date to use current date")
+    
+    start_date = get_date_input("Start date: ")
+    
+    # Default end date is current date
+    default_end = datetime.now()
+    end_date = get_date_input(f"End date [{default_end.strftime('%Y-%m-%d')}]: ", default_end)
+    
+    # Ensure end date is not before start date
+    if end_date < start_date:
+        print("End date cannot be before start date. Setting end date to start date.")
+        end_date = start_date
+        
+    return start_date, end_date
+
 def get_interval_choice() -> str:
     """Get user's preferred time interval"""
     print("\nSelect Time Interval:")