From d0b8e294dbafe22651e3e6853a5264ad4122dce1 Mon Sep 17 00:00:00 2001
From: "Bobby (aider)" <bobby@abellana.work>
Date: Mon, 17 Feb 2025 16:51:09 -0800
Subject: [PATCH] feat: Add shares field to watchlist items and update prefill
 functionality

---
 src/pages/trading/trading_system_page.py | 12 ++++++++++--
 src/trading/watchlist.py                 |  5 ++++-
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/src/pages/trading/trading_system_page.py b/src/pages/trading/trading_system_page.py
index f53284a..c11d524 100644
--- a/src/pages/trading/trading_system_page.py
+++ b/src/pages/trading/trading_system_page.py
@@ -202,7 +202,8 @@ def trading_system_page():
                     'ticker': ticker,
                     'entry_price': float(entry_price),
                     'target_price': float(target_price),
-                    'stop_loss': float(position['stop_loss'])
+                    'stop_loss': float(position['stop_loss']),
+                    'shares': recommended_shares
                 }
                 st.success("Details saved! Switch to Watch Lists tab to complete adding to your watch list.")
                 
@@ -247,6 +248,10 @@ def trading_system_page():
                                                value=float(prefill.get('entry_price', 0.0)),
                                                min_value=0.0,
                                                step=0.01)
+                    shares = st.number_input("Shares",
+                                           value=int(prefill.get('shares', 0)),
+                                           min_value=0,
+                                           step=1)
                 with col2:
                     target_price = st.number_input("Target Price",
                                                  value=float(prefill.get('target_price', 0.0)),
@@ -267,6 +272,7 @@ def trading_system_page():
                             entry_price=entry_price,
                             target_price=target_price,
                             stop_loss=stop_loss,
+                            shares=shares,
                             notes=notes
                         )
                         
@@ -302,7 +308,7 @@ def trading_system_page():
             if items:
                 for item in items:
                     with st.container():
-                        col1, col2, col3, col4, col5 = st.columns([2, 2, 2, 2, 1])
+                        col1, col2, col3, col4, col5, col6 = st.columns([2, 2, 2, 2, 1, 1])
                         with col1:
                             st.write(f"**{item.ticker}**")
                         with col2:
@@ -312,6 +318,8 @@ def trading_system_page():
                         with col4:
                             st.write(f"Stop: ${item.stop_loss:.2f}")
                         with col5:
+                            st.write(f"Shares: {item.shares:,}")
+                        with col6:
                             if st.button("Remove", key=f"remove_{item.id}"):
                                 if remove_from_watchlist(item.id):
                                     st.rerun()
diff --git a/src/trading/watchlist.py b/src/trading/watchlist.py
index 2dc93ee..9687a28 100644
--- a/src/trading/watchlist.py
+++ b/src/trading/watchlist.py
@@ -32,6 +32,7 @@ def ensure_tables_exist():
                 entry_price Float64,
                 target_price Float64,
                 stop_loss Float64,
+                shares Int32,
                 notes String,
                 created_at DateTime
             )
@@ -45,6 +46,7 @@ class WatchlistItem:
     entry_price: float
     target_price: float
     stop_loss: float
+    shares: Optional[int] = None
     notes: Optional[str] = None
     watchlist_id: Optional[int] = None
     id: Optional[int] = None
@@ -81,6 +83,7 @@ def add_to_watchlist(watchlist_id: int, item: WatchlistItem) -> bool:
                 item.entry_price,
                 item.target_price,
                 item.stop_loss,
+                item.shares or 0,
                 item.notes or '',
                 datetime.now()
             )]
@@ -90,7 +93,7 @@ def add_to_watchlist(watchlist_id: int, item: WatchlistItem) -> bool:
                 data,
                 column_names=[
                     'id', 'watchlist_id', 'ticker', 'entry_price', 
-                    'target_price', 'stop_loss', 'notes', 'created_at'
+                    'target_price', 'stop_loss', 'shares', 'notes', 'created_at'
                 ]
             )