利用DeepSeek辅助生成股市行情模拟数据测试电子表格插件rusty_sheet 0.2多表读取功能

因为要对DuckDB电子表格插件测试多表读取功能，所以让他帮助生成。
要求很简单，“编写python程序，生成每日一个sheet的xlsx文件，每个sheet1万行，是模拟的股市行情数据”，他也很明白数据结构，一举成功

我来为您编写一个生成模拟股市行情数据的Python程序，可以创建包含多个sheet的xlsx文件。

import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import random

def generate_stock_data(date, num_rows=10000):
    """
    生成指定日期的模拟股市行情数据
    
    Args:
        date: 日期 (datetime对象)
        num_rows: 数据行数
    
    Returns:
        DataFrame: 包含模拟股市数据的DataFrame
    """
    # 生成股票代码列表 (模拟1000只股票，每只股票10条记录)
    stock_codes = [f"STK{str(i).zfill(5)}" for i in range(1, 1001)]
    
    # 确保数据量足够
    if len(stock_codes) * 10 < num_rows:
        stock_codes = [f"STK{str(i).zfill(5)}" for i in range(1, num_rows//10 + 1)]
    
    data = []
    
    for i in range(num_rows):
        stock_code = random.choice(stock_codes)
        
        # 生成时间 (从9:30到15:00)
        hour = random.randint(9, 14)
        minute = random.randint(0, 59)
        if hour == 9:
            minute = random.randint(30, 59)  # 9:30开始
        elif hour == 14:
            minute = random.randint(0, 57)  # 15:00结束，留出3分钟缓冲
        
        timestamp = datetime(date.year, date.month, date.day, hour, minute, 
                            random.randint(0, 59))
        
        # 生成价格数据
        base_price = random.uniform(5, 500)  # 基础价格
        open_price = round(base_price, 2)
        high_price = round(open_price * random.uniform(1.0, 1.1), 2)
        low_price = round(open_price * random.uniform(0.9, 1.0), 2)
        close_price = round(random.uniform(low_price, high_price), 2)
        
        # 计算涨跌幅
        prev_close = round(open_price * random.uniform(0.95, 1.05), 2)
        change = close_price - prev_close
        change_percent = round((change / prev_close) * 100, 2)
        
        # 生成成交量数据
        volume = random.randint(10000, 10000000)
        amount = round(volume * close_price, 2)
        
        data.append({
            '股票代码': stock_code,
            '时间': timestamp,
            '开盘价': open_price,
            '最高价': high_price,
            '最低价': low_price,
            '收盘价': close_price,
            '前收盘价': prev_close,
            '涨跌额': round(change, 2),
            '涨跌幅(%)': change_percent,
            '成交量(股)': volume,
            '成交额(元)': amount,
            '换手率(%)': round(random.uniform(0.1, 15.0), 2)
        })
    
    df = pd.DataFrame(data)
    # 按时间排序
    df = df.sort_values('时间').reset_index(drop=True)
    
    return df

def create_daily_stock_xlsx(filename, start_date, num_days=30, rows_per_sheet=10000):
    """
    创建包含多个sheet的xlsx文件，每个sheet代表一天的股市数据
    
    Args:
        filename: 输出文件名
        start_date: 开始日期 (YYYY-MM-DD格式字符串或datetime对象)
        num_days: 生成的天数
        rows_per_sheet: 每个sheet的数据行数
    """
    if isinstance(start_date, str):
        start_date = datetime.strptime(start_date, '%Y-%m-%d')
    
    # 创建ExcelWriter对象
    with pd.ExcelWriter(filename, engine='openpyxl') as writer:
        
        for day in range(num_days):
            current_date = start_date + timedelta(days=day)
            sheet_name = current_date.strftime('%Y-%m-%d')
            
            print(f"正在生成 {sheet_name} 的数据...")
            
            # 生成当天的数据
            df = generate_stock_data(current_date, rows_per_sheet)
            
            # 写入到sheet
            df.to_excel(writer, sheet_name=sheet_name, index=False)
            
            # 调整列宽
            worksheet = writer.sheets[sheet_name]
            for idx, col in enumerate(df.columns):
                column_width = max(df[col].astype(str).map(len).max(), len(col)) + 2
                worksheet.column_dimensions[chr(65 + idx)].width = min(column_width, 20)
    
    print(f"\n文件 '{filename}' 生成完成！")
    print(f"共生成 {num_days} 天的数据，每个sheet包含 {rows_per_sheet} 行记录")

def main():
    """主函数"""
    # 配置参数
    output_filename = "daily_stock_data.xlsx"
    start_date = "2024-01-01"  # 开始日期
    num_days = 7               # 生成7天的数据
    rows_per_sheet = 10000     # 每个sheet的行数
    
    print("开始生成模拟股市行情数据...")
    print(f"文件名: {output_filename}")
    print(f"开始日期: {start_date}")
    print(f"生成天数: {num_days}")
    print(f"每页行数: {rows_per_sheet}")
    print("-" * 50)
    
    try:
        create_daily_stock_xlsx(
            filename=output_filename,
            start_date=start_date,
            num_days=num_days,
            rows_per_sheet=rows_per_sheet
        )
        
        # 显示生成数据的统计信息
        print("\n数据统计:")
        df_sample = pd.read_excel(output_filename, sheet_name=0)
        print(f"每行记录包含字段: {list(df_sample.columns)}")
        print(f"总记录数: {num_days * rows_per_sheet:,} 行")
        print(f"数据时间范围: {start_date} 到 {(datetime.strptime(start_date, '%Y-%m-%d') + timedelta(days=num_days-1)).strftime('%Y-%m-%d')}")
        
    except Exception as e:
        print(f"生成文件时出错: {e}")

if __name__ == "__main__":
    main()

这个程序还需要安装一些依赖包，请先安装：

pip install pandas openpyxl numpy

程序特点：

1. 数据结构完整：包含股票代码、时间、开盘价、最高价、最低价、收盘价、成交量、成交额、涨跌幅等完整字段
2. 数据真实性：
   • 交易时间限定在9:30-15:00
   • 价格波动符合实际规律
   • 成交量、成交额合理
3. 灵活的配置：
   • 可自定义开始日期
   • 可设置生成天数
   • 可调整每个sheet的行数
4. 自动格式化：
   • 自动调整列宽
   • 按时间排序
   • 日期格式的sheet名称

使用示例：

# 生成30天的数据，从2024年1月1日开始
create_daily_stock_xlsx("stock_data_2024.xlsx", "2024-01-01", 30, 10000)

# 生成最近7天的数据
create_daily_stock_xlsx("recent_stocks.xlsx", datetime.now().date() - timedelta(days=7), 7, 10000)

运行程序后，将生成一个xlsx文件，每个sheet代表一天的股市数据，包含1万行模拟的行情记录。数据格式规范，可以直接用于数据分析或测试。

测试记录如下

root@66d4e20ec1d7:/par# python3 stock.py
开始生成模拟股市行情数据...
文件名: daily_stock_data.xlsx
开始日期: 2024-01-01
生成天数: 7
每页行数: 10000
--------------------------------------------------
正在生成 2024-01-01 的数据...
正在生成 2024-01-02 的数据...
正在生成 2024-01-03 的数据...
正在生成 2024-01-04 的数据...
正在生成 2024-01-05 的数据...
正在生成 2024-01-06 的数据...
正在生成 2024-01-07 的数据...

文件 'daily_stock_data.xlsx' 生成完成！
共生成 7 天的数据，每个sheet包含 10000 行记录

数据统计:
每行记录包含字段: ['股票代码', '时间', '开盘价', '最高价', '最低价', '收盘价', '前收盘价', '涨跌额', '涨跌幅(%)', '成交量(股)', '成交额(元)', '换手率(%)']
总记录数: 70,000 行
数据时间范围: 2024-01-01 到 2024-01-07
root@66d4e20ec1d7:/par# ./duckdb140 -unsigned -cmd "load '/par/14/rusty_sheet.duckdb_extension';"
DuckDB v1.4.0 (Andium) b8a06e4a22
Enter ".help" for usage hints.
D .timer on
D copy (from read_sheets('daily_stock_data.xlsx')) to 'daily_stock_data.csv';
Run Time (s): real 1.118 user 1.232000 sys 0.084000
D copy (from read_sheet('daily_stock_data.xlsx',sheet_name='2024-01-01') union all from read_sheet('daily_stock_data.xlsx',sheet_name='2024-01-02')) to 'daily_stock_data2.csv';
Run Time (s): real 0.361 user 0.344000 sys 0.016000
D copy (from read_sheet('daily_stock_data.xlsx',sheet_name='2024-01-01') union all from read_sheet('daily_stock_data.xlsx',sheet_name='2024-01-02')union all from read_sheet('daily_stock_data.xlsx',sheet_name='2024-01-03')union all from read_sheet('daily_stock_data.xlsx',sheet_name='2024-01-04')) to 'daily_stock_data4.csv';
Run Time (s): real 0.581 user 0.620000 sys 0.020000

理论上，用 read_sheets 可以避免多个read_sheet重复解析和重复读取，效率更高。实测结果两者差距微小，主要还是自动处理更方便了。
补记，用更大的包含60个sheet的xlsx文件来测试，read_sheets效率提升还是很明显的。为节省篇幅，将union all语句写到脚本了。

D copy (from read_sheets('daily_stock_data.xlsx')) to 'daily_stock_data.csv';
Run Time (s): real 5.373 user 3.434030 sys 0.648087
D .read test_copy_sheet.txt
Run Time (s): real 15.273 user 4.458996 sys 0.863053
D .read test_copy_xlsx.txt
100% ▕██████████████████████████████████████▏ (00:00:09.63 elapsed)
Run Time (s): real 15.462 user 12.236893 sys 1.008438