时间表生成器

代码:

import sys
from datetime import datetime

import pandas as pd
from PIL import Image, ImageDraw, ImageFont
from PySide6.QtCore import QDate, QPoint, Qt
from PySide6.QtGui import QAction
from PySide6.QtWidgets import (
QAbstractItemView,
QApplication,
QDateEdit,
QFileDialog,
QHBoxLayout,
QHeaderView,
QLabel,
QLineEdit,
QMainWindow,
QMenu,
QMessageBox,
QPushButton,
QTableWidget,
QTableWidgetItem,
QVBoxLayout,
QWidget,
)


class TimelineGenerator(QMainWindow):
    def __init__(self):
        super().__init__()
        self.initUI()

    def initUI(self):
        self.setWindowTitle("时间表生成器")
        self.setGeometry(100, 100, 1200, 800)

        # 中央部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout(central_widget)

        # 标题设置
        title_layout = QHBoxLayout()
        title_layout.addWidget(QLabel("图表标题:"))
        self.title_input = QLineEdit("时间表")
        self.title_input.setMinimumWidth(300)
        title_layout.addWidget(self.title_input)
        title_layout.addStretch()
        layout.addLayout(title_layout)

        # 数据表格
        self.table = QTableWidget()
        self.table.setColumnCount(3)
        self.table.setHorizontalHeaderLabels(["名称", "开始日期", "结束日期"])
        self.table.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
        self.table.setContextMenuPolicy(Qt.CustomContextMenu)
        self.table.customContextMenuRequested.connect(self.show_context_menu)
        self.table.setSelectionBehavior(QAbstractItemView.SelectRows)

        # 初始化示例数据
        self.init_sample_data()
        layout.addWidget(self.table)

        # 生成图片按钮 - 占满整行
        self.generate_btn = QPushButton("生成图片")
        self.generate_btn.clicked.connect(self.generate_image)
        self.generate_btn.setMinimumHeight(40)
        layout.addWidget(self.generate_btn)

    def init_sample_data(self):
        """初始化数据"""
        sample_data = []

        self.table.setRowCount(len(sample_data))
        for i, (name, start, end) in enumerate(sample_data):
            self.table.setItem(i, 0, QTableWidgetItem(name))

            # 使用日期选择器
            start_date = QDateEdit()
            start_date.setDate(QDate.fromString(start, "yyyy-MM-dd"))
            start_date.setCalendarPopup(True)
            start_date.setDisplayFormat("yyyy-MM-dd")
            self.table.setCellWidget(i, 1, start_date)

            end_date = QDateEdit()
            end_date.setDate(QDate.fromString(end, "yyyy-MM-dd"))
            end_date.setCalendarPopup(True)
            end_date.setDisplayFormat("yyyy-MM-dd")
            self.table.setCellWidget(i, 2, end_date)

    def show_context_menu(self, position: QPoint):
        """显示右键菜单"""
        menu = QMenu(self)

        insert_above_action = QAction("在上方插入行", self)
        insert_above_action.triggered.connect(lambda: self.insert_row_above())
        menu.addAction(insert_above_action)

        insert_below_action = QAction("在下方插入行", self)
        insert_below_action.triggered.connect(lambda: self.insert_row_below())
        menu.addAction(insert_below_action)

        menu.addSeparator()

        delete_action = QAction("删除行", self)
        delete_action.triggered.connect(self.delete_row)
        menu.addAction(delete_action)

        menu.exec(self.table.viewport().mapToGlobal(position))

    def insert_row_above(self):
        """在选中行上方插入新行"""
        current_row = self.table.currentRow()
        if current_row < 0:
            current_row = 0
        self.insert_row(current_row)

    def insert_row_below(self):
        """在选中行下方插入新行"""
        current_row = self.table.currentRow()
        if current_row < 0:
            current_row = self.table.rowCount() - 1
        self.insert_row(current_row + 1)

    def insert_row(self, row: int):
        """在指定位置插入新行"""
        self.table.insertRow(row)

        # 设置默认日期
        today = QDate.currentDate()

        # 名称单元格
        self.table.setItem(row, 0, QTableWidgetItem(""))

        # 开始日期选择器
        start_date = QDateEdit()
        start_date.setDate(today)
        start_date.setCalendarPopup(True)
        start_date.setDisplayFormat("yyyy-MM-dd")
        self.table.setCellWidget(row, 1, start_date)

        # 结束日期选择器（默认比开始日期晚一天）
        end_date = QDateEdit()
        end_date.setDate(today.addDays(1))
        end_date.setCalendarPopup(True)
        end_date.setDisplayFormat("yyyy-MM-dd")
        self.table.setCellWidget(row, 2, end_date)

    def delete_row(self):
        """删除选中行"""
        current_row = self.table.currentRow()
        if current_row >= 0:
            self.table.removeRow(current_row)

    def get_data(self):
        """从表格获取数据"""
        data = []
        for row in range(self.table.rowCount()):
            name_item = self.table.item(row, 0)
            start_widget = self.table.cellWidget(row, 1)
            end_widget = self.table.cellWidget(row, 2)

            if (
                name_item is not None
                and start_widget is not None
                and end_widget is not None
            ):
                try:
                    name = name_item.text().strip()

                    # 从日期选择器获取日期并转换为datetime.datetime
                    start_qdate = start_widget.date()
                    end_qdate = end_widget.date()

                    start_date = datetime(
                        start_qdate.year(), start_qdate.month(), start_qdate.day()
                    )
                    end_date = datetime(
                        end_qdate.year(), end_qdate.month(), end_qdate.day()
                    )

                    display_name = name.strip()
                    data.append(
                        {
                            "name": name,
                            "display_name": display_name,
                            "start": start_date,
                            "end": end_date,
                        }
                    )
                except Exception as e:
                    QMessageBox.warning(
                        self, "数据格式错误", f"第{row+1}行数据错误: {e}"
                    )
                    return None
        return data

    def generate_image(self):
        """生成时间表图片"""
        data = self.get_data()
        if not data:
            QMessageBox.warning(self, "数据错误", "没有有效数据")
            return

        # 验证日期
        for i, item in enumerate(data):
            if item["end"] <= item["start"]:
                QMessageBox.warning(self, "日期错误", f"第{i+1}行：结束日期必须大于开始日期")
                return

        title = self.title_input.text().strip()
        if not title:
            title = "时间表"

        # 选择保存路径
        file_path, _ = QFileDialog.getSaveFileName(
            self, "保存图片", f"{title}.png", "PNG Images (*.png)"
        )

        if file_path:
            try:
                self.create_timeline_image(title, data, file_path)
                QMessageBox.information(self, "成功", f"图片已保存到: {file_path}")
            except Exception as e:
                QMessageBox.critical(self, "生成失败", f"生成图片时出错: {e}")

    def create_timeline_image(self, title, data, output_path):
        """创建时间表图片"""
        # 动态计算图片尺寸
        width = 1400
        height = 200 + len(data) * 70

        # 创建图片
        image = Image.new("RGB", (width, height), "white")
        draw = ImageDraw.Draw(image)

        # 字体设置
        try:
            title_font = ImageFont.truetype("simhei.ttf", 32)
            item_font = ImageFont.truetype("simhei.ttf", 20)
            date_font = ImageFont.truetype("simhei.ttf", 14)
            small_font = ImageFont.truetype("simhei.ttf", 12)
        except:
            title_font = ImageFont.load_default()
            item_font = ImageFont.load_default()
            date_font = ImageFont.load_default()
            small_font = ImageFont.load_default()

        # 绘制标题
        title_bbox = draw.textbbox((0, 0), title, font=title_font)
        title_width = title_bbox[2] - title_bbox[0]
        title_x = (width - title_width) // 2
        draw.text((title_x, 30), title, fill="black", font=title_font)

        # 计算时间范围
        all_starts = [item["start"] for item in data]
        all_ends = [item["end"] for item in data]
        min_date = min(all_starts)
        max_date = max(all_ends)
        date_range = max(1, (max_date - min_date).days)  # 确保至少为1

        # 图表参数
        chart_top = 100
        chart_bottom = height - 50
        chart_left = 250
        chart_right = width - 50
        chart_width = chart_right - chart_left
        row_height = 60

        # 绘制时间轴
        timeline_y = chart_bottom - 20
        draw.line(
            [(chart_left, timeline_y), (chart_right, timeline_y)], fill="black", width=2
        )

        # 绘制年份刻度
        start_year = min_date.year
        end_year = max_date.year

        for year in range(start_year, end_year + 2):
            year_date = datetime(year, 1, 1)
            if year_date < min_date:
                continue

            days_from_start = (year_date - min_date).days
            x_pos = chart_left + (days_from_start / date_range) * chart_width

            if chart_left <= x_pos <= chart_right:
                draw.line(
                    [(x_pos, chart_top), (x_pos, timeline_y)], fill="lightgray", width=1
                )
                year_str = f"{year}"
                year_bbox = draw.textbbox((0, 0), year_str, font=small_font)
                year_width = year_bbox[2] - year_bbox[0]
                draw.text(
                    (x_pos - year_width // 2, timeline_y + 10),
                    year_str,
                    fill="gray",
                    font=small_font,
                )

        # 绘制每个项目的时间条
        colors = [
            "#FF6B6B",
            "#4ECDC4",
            "#45B7D1",
            "#96CEB4",
            "#FFEAA7",
            "#DDA0DD",
            "#98D8C8",
            "#F7DC6F",
        ]

        for i, item in enumerate(data):
            row_y = chart_top + i * row_height

            # 计算开始和结束位置
            start_days = (item["start"] - min_date).days
            end_days = (item["end"] - min_date).days

            original_start_x = chart_left + (start_days / date_range) * chart_width
            original_end_x = chart_left + (end_days / date_range) * chart_width

            # 保存原始位置，用于日期标签的参考
            start_x = original_start_x
            end_x = original_end_x

            # 确保时间条至少有一定宽度
            min_bar_width = 15  # 最小时间条宽度
            if end_x - start_x < min_bar_width:
                # 如果时间太短，以中心点扩展
                center_x = (original_start_x + original_end_x) / 2
                start_x = max(chart_left, center_x - min_bar_width / 2)
                end_x = min(chart_right, center_x + min_bar_width / 2)

            # 绘制时间条
            bar_height = 25
            bar_y = row_y - bar_height // 2

            color = colors[i % len(colors)]
            draw.rounded_rectangle(
                [start_x, bar_y, end_x, bar_y + bar_height],
                radius=3,
                fill=color,
                outline="darkgray",
                width=1,
            )

            # 绘制项目名称
            display_name = item["display_name"]
            name_bbox = draw.textbbox((0, 0), display_name, font=item_font)
            draw.text(
                (50, row_y - name_bbox[3] // 2),
                display_name,
                fill="black",
                font=item_font,
            )

            # 绘制日期标签
            start_str = item["start"].strftime("%Y-%m-%d")
            end_str = item["end"].strftime("%Y-%m-%d")

            # 获取日期文本的尺寸
            start_bbox = draw.textbbox((0, 0), start_str, font=date_font)
            end_bbox = draw.textbbox((0, 0), end_str, font=date_font)
            start_text_width = start_bbox[2] - start_bbox[0]
            end_text_width = end_bbox[2] - end_bbox[0]

            # 检查日期是否过于接近（使用原始位置判断，避免被调整后的位置误导）
            date_spacing = (
                (original_end_x - original_start_x) / chart_width * date_range
            )

            if date_spacing < 30:  # 如果日期间隔小于30天
                # 只显示一个合并的日期标签
                # 使用调整后的时间条中心点
                mid_x = (start_x + end_x) / 2
                date_str = f"{start_str} ~ {end_str}"
                date_bbox = draw.textbbox((0, 0), date_str, font=date_font)
                date_width = date_bbox[2] - date_bbox[0]

                # 确保合并标签在可视范围内
                date_x = mid_x - date_width / 2
                date_x = max(chart_left + 5, min(date_x, chart_right - date_width - 5))

                draw.text(
                    (date_x, bar_y + bar_height + 5),
                    date_str,
                    fill="darkblue",
                    font=date_font,
                )

                # 移除有style参数的虚线，改为实线
                draw.line(
                    [
                        (original_start_x, bar_y + bar_height),
                        (original_start_x, bar_y + bar_height + 15),
                    ],
                    fill="gray",
                    width=1,
                )
                draw.line(
                    [
                        (original_end_x, bar_y + bar_height),
                        (original_end_x, bar_y + bar_height + 15),
                    ],
                    fill="gray",
                    width=1,
                )
            else:
                # 分别显示开始和结束日期
                # 对于被调整过的时间条，日期标签应该对应原始位置
                label_start_x = original_start_x
                label_end_x = original_end_x

                # 开始日期
                start_text_x = label_start_x - start_text_width / 2

                # 检查是否超出左边界
                if start_text_x < chart_left + 5:
                    start_text_x = chart_left + 5
                # 检查是否与结束日期重叠
                elif (
                    start_text_x + start_text_width > label_end_x - end_text_width - 10
                ):
                    start_text_x = max(chart_left + 5, label_start_x - start_text_width)

                draw.text(
                    (start_text_x, bar_y + bar_height + 5),
                    start_str,
                    fill="darkgreen",
                    font=date_font,
                )

                # 结束日期
                end_text_x = label_end_x - end_text_width / 2

                # 检查是否超出右边界
                if end_text_x + end_text_width > chart_right - 5:
                    end_text_x = chart_right - end_text_width - 5
                # 检查是否与开始日期重叠
                elif end_text_x < start_text_x + start_text_width + 10:
                    end_text_x = min(chart_right - end_text_width - 5, label_end_x)

                draw.text(
                    (end_text_x, bar_y + bar_height + 5),
                    end_str,
                    fill="darkred",
                    font=date_font,
                )

            # 绘制连接线到时间条
            # 使用原始位置画连接线，更准确地反映实际日期位置
            draw.line(
                [
                    (original_start_x, bar_y + bar_height),
                    (original_start_x, bar_y + bar_height + 5),
                ],
                fill="gray",
                width=1,
            )
            draw.line(
                [
                    (original_end_x, bar_y + bar_height),
                    (original_end_x, bar_y + bar_height + 5),
                ],
                fill="gray",
                width=1,
            )

        # 保存图片
        image.save(output_path)


def main():
    app = QApplication(sys.argv)
    window = TimelineGenerator()
    window.show()
    sys.exit(app.exec())


if __name__ == "__main__":
    main()