Files
2026-05-26 11:04:27 -04:00

356 lines
14 KiB
Python

from __future__ import annotations
import numpy as np
import pandas as pd
import pyqtgraph as pg
from PyQt6.QtWidgets import (
QWidget, QVBoxLayout, QHBoxLayout, QLabel, QLineEdit,
QPushButton, QCheckBox, QCompleter, QSizePolicy, QFrame
)
from PyQt6.QtCore import Qt, QStringListModel, QThread, pyqtSignal, QTimer
from PyQt6.QtGui import QPicture, QPainter, QColor, QPen, QBrush, QFont
from PyQt6.QtCore import QRectF, QPointF
from utils.formatters import fmt_currency, fmt_percent, pnl_color
PERIODS = ["1D", "1W", "1M", "3M", "6M", "1Y", "5Y"]
INDICATORS = ["SMA20", "SMA50", "EMA20", "BB", "Volume", "RSI", "MACD"]
COLORS = {
"up": "#a6e3a1",
"down": "#f38ba8",
"sma20": "#89b4fa",
"sma50": "#fab387",
"ema20": "#cba6f7",
"bb_upper": "#89dceb",
"bb_lower": "#89dceb",
"bb_fill": "#1a3a4a",
"volume_up": "#a6e3a1",
"volume_down": "#f38ba8",
"rsi_line": "#cba6f7",
"macd_line": "#89b4fa",
"macd_signal": "#fab387",
"macd_pos": "#a6e3a1",
"macd_neg": "#f38ba8",
}
class CandlestickItem(pg.GraphicsObject):
def __init__(self, df: pd.DataFrame):
super().__init__()
self._df = df
self._picture = QPicture()
self._generate()
def _generate(self):
self._picture = QPicture()
p = QPainter(self._picture)
p.setRenderHint(QPainter.RenderHint.Antialiasing, False)
w = 0.3
for i, (ts, row) in enumerate(self._df.iterrows()):
t = float(i)
o, c, lo, hi = float(row["Open"]), float(row["Close"]), float(row["Low"]), float(row["High"])
is_up = c >= o
color = QColor(COLORS["up"] if is_up else COLORS["down"])
pen = QPen(color)
pen.setWidthF(1.0)
p.setPen(pen)
p.setBrush(QBrush(color))
p.drawLine(QPointF(t, lo), QPointF(t, hi))
body_top = max(o, c)
body_bot = min(o, c)
body_h = max(body_top - body_bot, 0.001)
p.drawRect(QRectF(t - w, body_bot, w * 2, body_h))
p.end()
def paint(self, p, *args):
p.drawPicture(0, 0, self._picture)
def boundingRect(self):
return QRectF(self._picture.boundingRect())
class DateAxisItem(pg.AxisItem):
def __init__(self, timestamps: list, *args, **kwargs):
super().__init__(*args, **kwargs)
self._timestamps = timestamps
def tickStrings(self, values, scale, spacing):
result = []
for v in values:
idx = int(round(v))
if 0 <= idx < len(self._timestamps):
ts = self._timestamps[idx]
if hasattr(ts, "strftime"):
result.append(ts.strftime("%m/%d"))
else:
result.append(str(ts)[:10])
else:
result.append("")
return result
class ChartDataWorker(QThread):
data_ready = pyqtSignal(object, str)
error = pyqtSignal(str)
def __init__(self, symbol: str, period: str):
super().__init__()
self.symbol = symbol
self.period = period
def run(self):
try:
from core.market_data import get_chart_data
df = get_chart_data(self.symbol, self.period)
self.data_ready.emit(df, self.symbol)
except Exception as e:
self.error.emit(str(e))
class ChartWidget(QWidget):
symbol_changed = pyqtSignal(str)
def __init__(self, parent=None):
super().__init__(parent)
self._symbol = "AAPL"
self._period = "6M"
self._df: pd.DataFrame | None = None
self._worker: ChartDataWorker | None = None
self._active_indicators: set[str] = {"Volume"}
self._setup_ui()
self._setup_autocomplete()
self.load_symbol(self._symbol)
def _setup_ui(self):
root = QVBoxLayout(self)
root.setContentsMargins(12, 12, 12, 12)
root.setSpacing(8)
# Search + quote bar
top_bar = QHBoxLayout()
self._search = QLineEdit()
self._search.setPlaceholderText("Search symbol… (e.g. AAPL, BTC-USD)")
self._search.setMaximumWidth(280)
self._search.returnPressed.connect(self._on_search)
top_bar.addWidget(self._search)
self._quote_label = QLabel()
self._quote_label.setFont(QFont("Segoe UI", 11, QFont.Weight.Bold))
top_bar.addWidget(self._quote_label)
self._change_label = QLabel()
top_bar.addWidget(self._change_label)
top_bar.addStretch()
root.addLayout(top_bar)
# Period buttons
period_bar = QHBoxLayout()
self._period_btns: dict[str, QPushButton] = {}
for p in PERIODS:
btn = QPushButton(p)
btn.setCheckable(True)
btn.setMaximumWidth(48)
btn.clicked.connect(lambda checked, period=p: self._set_period(period))
self._period_btns[p] = btn
period_bar.addWidget(btn)
period_bar.addSpacing(20)
for ind in INDICATORS:
cb = QCheckBox(ind)
cb.setChecked(ind == "Volume")
cb.stateChanged.connect(lambda state, i=ind: self._toggle_indicator(i, state))
period_bar.addWidget(cb)
period_bar.addStretch()
root.addLayout(period_bar)
self._period_btns["6M"].setChecked(True)
# Chart area
pg.setConfigOptions(antialias=True, foreground="#6c7086", background="#1e1e2e")
self._layout_widget = pg.GraphicsLayoutWidget()
root.addWidget(self._layout_widget)
self._status_label = QLabel("Loading…")
self._status_label.setObjectName("subtitle")
root.addWidget(self._status_label)
self._build_plots()
def _build_plots(self):
self._layout_widget.clear()
self._main_plot: pg.PlotItem = self._layout_widget.addPlot(row=0, col=0)
self._main_plot.showGrid(x=False, y=True, alpha=0.15)
self._main_plot.getAxis("left").setWidth(60)
self._main_plot.getAxis("bottom").setStyle(showValues=False)
self._main_plot.setMinimumHeight(260)
self._vol_plot: pg.PlotItem = self._layout_widget.addPlot(row=1, col=0)
self._vol_plot.setMaximumHeight(80)
self._vol_plot.showGrid(x=False, y=True, alpha=0.1)
self._vol_plot.getAxis("left").setWidth(60)
self._vol_plot.setXLink(self._main_plot)
self._rsi_plot: pg.PlotItem = self._layout_widget.addPlot(row=2, col=0)
self._rsi_plot.setMaximumHeight(100)
self._rsi_plot.showGrid(x=False, y=True, alpha=0.1)
self._rsi_plot.getAxis("left").setWidth(60)
self._rsi_plot.setXLink(self._main_plot)
self._rsi_plot.setYRange(0, 100)
self._rsi_plot.hide()
self._macd_plot: pg.PlotItem = self._layout_widget.addPlot(row=3, col=0)
self._macd_plot.setMaximumHeight(100)
self._macd_plot.showGrid(x=False, y=True, alpha=0.1)
self._macd_plot.getAxis("left").setWidth(60)
self._macd_plot.setXLink(self._main_plot)
self._macd_plot.hide()
def _setup_autocomplete(self):
common = [
"AAPL", "MSFT", "GOOGL", "AMZN", "NVDA", "META", "TSLA", "SPY", "QQQ",
"JPM", "BAC", "V", "MA", "NFLX", "AMD", "INTC", "DIS", "WMT", "KO",
"BTC-USD", "ETH-USD", "SOL-USD", "BNB-USD", "XRP-USD",
]
self._completer_model = QStringListModel(common)
completer = QCompleter(self._completer_model)
completer.setCaseSensitivity(Qt.CaseSensitivity.CaseInsensitive)
self._search.setCompleter(completer)
def _on_search(self):
text = self._search.text().strip().upper()
if text:
self.load_symbol(text)
def _set_period(self, period: str):
for p, btn in self._period_btns.items():
btn.setChecked(p == period)
self._period = period
if self._symbol:
self.load_symbol(self._symbol)
def _toggle_indicator(self, indicator: str, state: int):
if state:
self._active_indicators.add(indicator)
else:
self._active_indicators.discard(indicator)
if self._df is not None:
self._render_chart(self._df)
def load_symbol(self, symbol: str):
self._symbol = symbol.upper()
self._status_label.setText(f"Loading {self._symbol}…")
self._quote_label.setText(self._symbol)
self._change_label.clear()
if self._worker and self._worker.isRunning():
self._worker.quit()
self._worker = ChartDataWorker(self._symbol, self._period)
self._worker.data_ready.connect(self._on_data_ready)
self._worker.error.connect(lambda e: self._status_label.setText(f"Error: {e}"))
self._worker.start()
self.symbol_changed.emit(self._symbol)
def _on_data_ready(self, df: pd.DataFrame, symbol: str):
if df.empty:
self._status_label.setText(f"No data for {symbol}")
return
self._df = df
self._render_chart(df)
self._status_label.setText(f"{symbol} | {len(df)} bars | {self._period}")
self._update_quote_display(df)
def _update_quote_display(self, df: pd.DataFrame):
last = float(df["Close"].iloc[-1])
prev = float(df["Close"].iloc[-2]) if len(df) > 1 else last
change = last - prev
change_pct = (change / prev * 100) if prev else 0
self._quote_label.setText(f"{self._symbol} {fmt_currency(last)}")
self._change_label.setText(f"{'+' if change >= 0 else ''}{fmt_currency(change)} {fmt_percent(change_pct, plus=True)}")
self._change_label.setStyleSheet(f"color: {pnl_color(change)};")
def _render_chart(self, df: pd.DataFrame):
self._build_plots()
xs = np.arange(len(df))
timestamps = list(df.index)
date_axis = DateAxisItem(timestamps, orientation="bottom")
self._main_plot.getAxis("bottom").hide()
candle = CandlestickItem(df)
self._main_plot.addItem(candle)
# Overlaid indicators
closes = df["Close"].values.astype(float)
if "SMA20" in self._active_indicators and len(closes) >= 20:
sma20 = pd.Series(closes).rolling(20).mean().values
self._main_plot.plot(xs, sma20, pen=pg.mkPen(COLORS["sma20"], width=1.5), name="SMA20")
if "SMA50" in self._active_indicators and len(closes) >= 50:
sma50 = pd.Series(closes).rolling(50).mean().values
self._main_plot.plot(xs, sma50, pen=pg.mkPen(COLORS["sma50"], width=1.5), name="SMA50")
if "EMA20" in self._active_indicators and len(closes) >= 20:
ema20 = pd.Series(closes).ewm(span=20).mean().values
self._main_plot.plot(xs, ema20, pen=pg.mkPen(COLORS["ema20"], width=1.5), name="EMA20")
if "BB" in self._active_indicators and len(closes) >= 20:
sma = pd.Series(closes).rolling(20).mean()
std = pd.Series(closes).rolling(20).std()
upper = (sma + 2 * std).values
lower = (sma - 2 * std).values
self._main_plot.plot(xs, upper, pen=pg.mkPen(COLORS["bb_upper"], width=1, style=Qt.PenStyle.DashLine))
self._main_plot.plot(xs, lower, pen=pg.mkPen(COLORS["bb_lower"], width=1, style=Qt.PenStyle.DashLine))
fill = pg.FillBetweenItem(
self._main_plot.plot(xs, upper, pen=None),
self._main_plot.plot(xs, lower, pen=None),
brush=pg.mkBrush(QColor(137, 220, 235, 30)),
)
self._main_plot.addItem(fill)
# Volume
if "Volume" in self._active_indicators:
volumes = df["Volume"].values.astype(float)
colors = [COLORS["volume_up"] if df["Close"].iloc[i] >= df["Open"].iloc[i] else COLORS["volume_down"] for i in range(len(df))]
brushes = [pg.mkBrush(QColor(c)) for c in colors]
bar = pg.BarGraphItem(x=xs, height=volumes, width=0.6, brushes=brushes, pens=[pg.mkPen(None)] * len(xs))
self._vol_plot.addItem(bar)
self._vol_plot.show()
else:
self._vol_plot.hide()
# RSI
if "RSI" in self._active_indicators and len(closes) >= 14:
delta = pd.Series(closes).diff()
gain = delta.clip(lower=0).rolling(14).mean()
loss = (-delta.clip(upper=0)).rolling(14).mean()
rs = gain / loss
rsi = (100 - (100 / (1 + rs))).values
self._rsi_plot.clear()
self._rsi_plot.plot(xs, rsi, pen=pg.mkPen(COLORS["rsi_line"], width=1.5))
self._rsi_plot.addLine(y=70, pen=pg.mkPen("#f38ba8", width=1, style=Qt.PenStyle.DashLine))
self._rsi_plot.addLine(y=30, pen=pg.mkPen("#a6e3a1", width=1, style=Qt.PenStyle.DashLine))
self._rsi_plot.show()
else:
self._rsi_plot.hide()
# MACD
if "MACD" in self._active_indicators and len(closes) >= 26:
ema12 = pd.Series(closes).ewm(span=12).mean()
ema26 = pd.Series(closes).ewm(span=26).mean()
macd_line = (ema12 - ema26).values
signal_line = pd.Series(macd_line).ewm(span=9).mean().values
histogram = macd_line - signal_line
self._macd_plot.clear()
hist_colors = [COLORS["macd_pos"] if v >= 0 else COLORS["macd_neg"] for v in histogram]
bar2 = pg.BarGraphItem(x=xs, height=histogram, width=0.6,
brushes=[pg.mkBrush(QColor(c)) for c in hist_colors],
pens=[pg.mkPen(None)] * len(xs))
self._macd_plot.addItem(bar2)
self._macd_plot.plot(xs, macd_line, pen=pg.mkPen(COLORS["macd_line"], width=1.5))
self._macd_plot.plot(xs, signal_line, pen=pg.mkPen(COLORS["macd_signal"], width=1.5))
self._macd_plot.show()
else:
self._macd_plot.hide()