Files
StockMind/stockmind/ui/watchlist_widget.py
T
2026-05-26 11:04:27 -04:00

198 lines
7.5 KiB
Python

from __future__ import annotations
from PyQt6.QtWidgets import (
QWidget, QVBoxLayout, QHBoxLayout, QLabel, QLineEdit, QPushButton,
QTreeWidget, QTreeWidgetItem, QMenu, QInputDialog, QMessageBox, QComboBox
)
from PyQt6.QtCore import Qt, QThread, pyqtSignal, QTimer
from PyQt6.QtGui import QColor, QFont
from utils.formatters import fmt_currency, fmt_percent, pnl_color
class WatchlistFetchWorker(QThread):
quotes_ready = pyqtSignal(dict)
def __init__(self, symbols: list[str]):
super().__init__()
self._symbols = symbols
def run(self):
from core.market_data import get_batch_quotes
self.quotes_ready.emit(get_batch_quotes(self._symbols))
class WatchlistWidget(QWidget):
symbol_selected = pyqtSignal(str)
def __init__(self, parent=None):
super().__init__(parent)
self._worker: WatchlistFetchWorker | None = None
self._setup_ui()
self._load_from_db()
self._start_timer()
def _setup_ui(self):
root = QVBoxLayout(self)
root.setContentsMargins(12, 12, 12, 12)
root.setSpacing(8)
header = QLabel("Watchlist")
header.setObjectName("title")
root.addWidget(header)
# Add symbol row
add_row = QHBoxLayout()
self._symbol_input = QLineEdit()
self._symbol_input.setPlaceholderText("Add symbol…")
self._symbol_input.returnPressed.connect(self._add_symbol)
self._group_combo = QComboBox()
self._group_combo.setEditable(True)
self._group_combo.addItems(["Default", "Tech", "Crypto", "ETFs"])
self._group_combo.setMaximumWidth(100)
add_btn = QPushButton("Add")
add_btn.setObjectName("primary_btn")
add_btn.clicked.connect(self._add_symbol)
add_row.addWidget(self._symbol_input)
add_row.addWidget(self._group_combo)
add_row.addWidget(add_btn)
root.addLayout(add_row)
# Tree
self._tree = QTreeWidget()
self._tree.setColumnCount(4)
self._tree.setHeaderLabels(["Symbol", "Price", "Change", "% Change"])
self._tree.header().setStretchLastSection(False)
self._tree.setColumnWidth(0, 120)
self._tree.setColumnWidth(1, 90)
self._tree.setColumnWidth(2, 80)
self._tree.setColumnWidth(3, 80)
self._tree.setAlternatingRowColors(True)
self._tree.itemDoubleClicked.connect(self._on_item_double_clicked)
self._tree.setContextMenuPolicy(Qt.ContextMenuPolicy.CustomContextMenu)
self._tree.customContextMenuRequested.connect(self._show_context_menu)
root.addWidget(self._tree)
self._status_label = QLabel("")
self._status_label.setObjectName("subtitle")
root.addWidget(self._status_label)
def _start_timer(self):
self._timer = QTimer(self)
self._timer.timeout.connect(self._refresh_quotes)
self._timer.start(60_000)
def _load_from_db(self):
from db.database import get_session
from db.models import WatchlistItem
self._tree.clear()
groups: dict[str, QTreeWidgetItem] = {}
with get_session() as session:
items = session.query(WatchlistItem).order_by(WatchlistItem.group_name).all()
for item in items:
group = item.group_name or "Default"
if group not in groups:
group_item = QTreeWidgetItem([group])
group_item.setFont(0, QFont("Segoe UI", 9, QFont.Weight.Bold))
group_item.setForeground(0, QColor("#89b4fa"))
self._tree.addTopLevelItem(group_item)
group_item.setExpanded(True)
groups[group] = group_item
child = QTreeWidgetItem([item.symbol, "—", "—", "—"])
child.setData(0, Qt.ItemDataRole.UserRole, item.symbol)
groups[group].addChild(child)
self._refresh_quotes()
def _refresh_quotes(self):
symbols = self._collect_symbols()
if not symbols or (self._worker and self._worker.isRunning()):
return
self._worker = WatchlistFetchWorker(symbols)
self._worker.quotes_ready.connect(self._apply_quotes)
self._worker.start()
def _collect_symbols(self) -> list[str]:
symbols = []
for i in range(self._tree.topLevelItemCount()):
group = self._tree.topLevelItem(i)
for j in range(group.childCount()):
sym = group.child(j).data(0, Qt.ItemDataRole.UserRole)
if sym:
symbols.append(sym)
return symbols
def _apply_quotes(self, quotes: dict):
from datetime import datetime
for i in range(self._tree.topLevelItemCount()):
group = self._tree.topLevelItem(i)
for j in range(group.childCount()):
child = group.child(j)
sym = child.data(0, Qt.ItemDataRole.UserRole)
q = quotes.get(sym, {})
price = q.get("price", 0)
change = q.get("change", 0)
change_pct = q.get("change_pct", 0)
child.setText(1, fmt_currency(price))
child.setText(2, f"{'+' if change >= 0 else ''}{fmt_currency(change)}")
child.setText(3, fmt_percent(change_pct, plus=True))
color = QColor(pnl_color(change))
for col in (1, 2, 3):
child.setForeground(col, color)
self._status_label.setText(f"Updated {datetime.now().strftime('%H:%M:%S')}")
def _add_symbol(self):
symbol = self._symbol_input.text().strip().upper()
group = self._group_combo.currentText().strip() or "Default"
if not symbol:
return
from db.database import get_session
from db.models import WatchlistItem
with get_session() as session:
existing = session.query(WatchlistItem).filter_by(symbol=symbol).first()
if existing:
return
session.add(WatchlistItem(symbol=symbol, group_name=group))
session.commit()
self._symbol_input.clear()
self._load_from_db()
if self._group_combo.findText(group) == -1:
self._group_combo.addItem(group)
def _on_item_double_clicked(self, item: QTreeWidgetItem, column: int):
sym = item.data(0, Qt.ItemDataRole.UserRole)
if sym:
self.symbol_selected.emit(sym)
def _show_context_menu(self, pos):
item = self._tree.itemAt(pos)
if not item or not item.data(0, Qt.ItemDataRole.UserRole):
return
sym = item.data(0, Qt.ItemDataRole.UserRole)
menu = QMenu(self)
view_action = menu.addAction(f"View Chart: {sym}")
remove_action = menu.addAction("Remove")
action = menu.exec(self._tree.mapToGlobal(pos))
if action == view_action:
self.symbol_selected.emit(sym)
elif action == remove_action:
self._remove_symbol(sym)
def _remove_symbol(self, symbol: str):
from db.database import get_session
from db.models import WatchlistItem
with get_session() as session:
session.query(WatchlistItem).filter_by(symbol=symbol).delete()
session.commit()
self._load_from_db()
def add_symbol_external(self, symbol: str, group: str = "Default"):
self._symbol_input.setText(symbol)
idx = self._group_combo.findText(group)
if idx >= 0:
self._group_combo.setCurrentIndex(idx)
self._add_symbol()