From 0ce4bb448600b85e7e90b5ce8720f81c8ecb201e Mon Sep 17 00:00:00 2001
From: "Bobby (aider)" <bobby@abellana.work>
Date: Mon, 10 Feb 2025 09:45:29 -0800
Subject: [PATCH] feat: Add timezone-aware datetime input with market hours
 validation

---
 src/trading/journal.py | 79 +++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 78 insertions(+), 1 deletion(-)

diff --git a/src/trading/journal.py b/src/trading/journal.py
index 12c6f5f..979145b 100644
--- a/src/trading/journal.py
+++ b/src/trading/journal.py
@@ -1,6 +1,8 @@
-from datetime import datetime
+from datetime import datetime, timedelta
 from dataclasses import dataclass
 from typing import Optional
+import pytz
+from zoneinfo import ZoneInfo
 from db.db_connection import create_client
 from trading.position_calculator import PositionCalculator
 
@@ -36,6 +38,74 @@ class TradeEntry:
             return (self.exit_price - self.entry_price) * self.shares
         return None
 
+def get_market_hours(date: datetime) -> tuple:
+    """Get market open/close times in Eastern for given date"""
+    eastern = pytz.timezone('US/Eastern')
+    date_eastern = date.astimezone(eastern)
+    
+    market_open = eastern.localize(
+        datetime.combine(date_eastern.date(), datetime.strptime("09:30", "%H:%M").time())
+    )
+    market_close = eastern.localize(
+        datetime.combine(date_eastern.date(), datetime.strptime("16:00", "%H:%M").time())
+    )
+    return market_open, market_close
+
+def validate_market_time(dt: datetime) -> tuple[datetime, bool]:
+    """
+    Validate if time is during market hours, adjust if needed
+    Returns: (adjusted_datetime, was_adjusted)
+    """
+    pacific = pytz.timezone('US/Pacific')
+    eastern = pytz.timezone('US/Eastern')
+    
+    # Ensure datetime is timezone-aware
+    if dt.tzinfo is None:
+        dt = pacific.localize(dt)
+    
+    dt_eastern = dt.astimezone(eastern)
+    market_open, market_close = get_market_hours(dt_eastern)
+    
+    if dt_eastern < market_open:
+        return market_open.astimezone(pacific), True
+    elif dt_eastern > market_close:
+        return market_close.astimezone(pacific), True
+    
+    return dt, False
+
+def get_datetime_input(prompt: str, default: datetime = None) -> datetime:
+    """Get date and time input in Pacific time"""
+    pacific = pytz.timezone('US/Pacific')
+    
+    while True:
+        try:
+            if default:
+                print(f"Press Enter for current time ({default.strftime('%Y-%m-%d %H:%M')})")
+            date_str = input(f"{prompt} (YYYY-MM-DD HH:MM): ").strip()
+            
+            if not date_str and default:
+                dt = default
+            else:
+                dt = datetime.strptime(date_str, "%Y-%m-%d %H:%M")
+            
+            # Make datetime timezone-aware (Pacific)
+            dt = pacific.localize(dt)
+            
+            # Validate market hours
+            adjusted_dt, was_adjusted = validate_market_time(dt)
+            if was_adjusted:
+                print(f"\nWarning: Time adjusted to market hours (Eastern)")
+                print(f"Original (Pacific): {dt.strftime('%Y-%m-%d %H:%M %Z')}")
+                print(f"Adjusted (Pacific): {adjusted_dt.strftime('%Y-%m-%d %H:%M %Z')}")
+                print(f"Adjusted (Eastern): {adjusted_dt.astimezone(pytz.timezone('US/Eastern')).strftime('%Y-%m-%d %H:%M %Z')}")
+                if input("Accept adjusted time? (y/n): ").lower() != 'y':
+                    continue
+            
+            return adjusted_dt
+            
+        except ValueError:
+            print("Invalid format. Please use YYYY-MM-D HH:MM")
+
 def create_trades_table():
     with create_client() as client:
         query = """
@@ -208,6 +278,9 @@ def journal_menu():
             else:
                 position_id = generate_position_id(ticker)
             
+            # Get entry date/time with market hours validation
+            entry_date = get_datetime_input("Enter entry date and time", default=datetime.now())
+            
             shares = int(input("Enter number of shares: "))
             entry_price = float(input("Enter entry price: "))
             order_type = get_order_type()
@@ -270,6 +343,10 @@ def journal_menu():
             
             trade_id = int(input("\nEnter trade ID to update: "))
             exit_price = float(input("Enter exit price: "))
+            
+            # Get exit date/time with market hours validation
+            exit_date = get_datetime_input("Enter exit date and time", default=datetime.now())
+            
             followed_rules = input("Did you follow your rules? (y/n): ").lower() == 'y'
             exit_reason = input("Enter exit reason: ")
             notes = input("Additional notes (optional): ") or None