EU-Utility/premium/widgets/chart_widget.py

146 lines
4.3 KiB
Python

"""
Chart widget for displaying data visualizations.
"""
from enum import Enum
from typing import List, Dict, Any, Optional
try:
from PyQt6.QtWidgets import QWidget, QVBoxLayout
from PyQt6.QtCore import Qt
from PyQt6.QtGui import QPainter, QColor, QPen, QBrush, QFont
HAS_QT = True
except ImportError:
HAS_QT = False
QWidget = object
class ChartType(Enum):
LINE = "line"
BAR = "bar"
PIE = "pie"
class ChartWidget(QWidget if HAS_QT else object):
"""Simple chart widget."""
def __init__(
self,
chart_type: ChartType = ChartType.LINE,
title: str = "Chart",
parent=None
):
if not HAS_QT:
return
super().__init__(parent)
self.chart_type = chart_type
self.title = title
self._data: List[Dict[str, Any]] = []
self._max_points = 50
self._y_max = 100
self.setMinimumHeight(150)
self.setStyleSheet("background-color: #2d2d2d; border-radius: 8px;")
def add_point(self, x: Any, y: float, label: Optional[str] = None):
"""Add a data point."""
self._data.append({'x': x, 'y': y, 'label': label})
# Limit points
if len(self._data) > self._max_points:
self._data = self._data[-self._max_points:]
# Update scale
if y > self._y_max:
self._y_max = y * 1.1
if HAS_QT:
self.update()
def set_data(self, data: List[Dict[str, Any]]):
"""Set all data points."""
self._data = data
if data:
self._y_max = max(d['y'] for d in data) * 1.1
if HAS_QT:
self.update()
def clear(self):
"""Clear all data."""
self._data = []
if HAS_QT:
self.update()
def paintEvent(self, event):
if not HAS_QT or not self._data:
return
painter = QPainter(self)
painter.setRenderHint(QPainter.RenderHint.Antialiasing)
width = self.width()
height = self.height()
padding = 40
chart_width = width - padding * 2
chart_height = height - padding * 2
# Background
painter.fillRect(self.rect(), QColor(45, 45, 45))
# Draw title
painter.setPen(QColor(255, 255, 255))
painter.setFont(QFont("Arial", 10, QFont.Weight.Bold))
painter.drawText(10, 20, self.title)
if len(self._data) < 2:
return
# Draw chart based on type
if self.chart_type == ChartType.LINE:
self._draw_line_chart(painter, padding, chart_width, chart_height)
elif self.chart_type == ChartType.BAR:
self._draw_bar_chart(painter, padding, chart_width, chart_height)
def _draw_line_chart(self, painter, padding, width, height):
"""Draw a line chart."""
if len(self._data) < 2:
return
pen = QPen(QColor(76, 175, 80))
pen.setWidth(2)
painter.setPen(pen)
x_step = width / max(len(self._data) - 1, 1)
# Draw line
points = []
for i, point in enumerate(self._data):
x = padding + i * x_step
y = padding + height - (point['y'] / self._y_max * height)
points.append((x, y))
for i in range(len(points) - 1):
painter.drawLine(int(points[i][0]), int(points[i][1]),
int(points[i+1][0]), int(points[i+1][1]))
# Draw points
painter.setBrush(QColor(76, 175, 80))
for x, y in points:
painter.drawEllipse(int(x - 3), int(y - 3), 6, 6)
def _draw_bar_chart(self, painter, padding, width, height):
"""Draw a bar chart."""
bar_width = width / len(self._data) * 0.8
gap = width / len(self._data) * 0.2
painter.setBrush(QColor(76, 175, 80))
painter.setPen(Qt.PenStyle.NoPen)
for i, point in enumerate(self._data):
bar_height = point['y'] / self._y_max * height
x = padding + i * (bar_width + gap) + gap / 2
y = padding + height - bar_height
painter.drawRect(int(x), int(y), int(bar_width), int(bar_height))