开源自动化数据分析工具——data-copilot源码分析

本文对开源数据分析工具data-copilot源代码中的关键部分，create_tool/Atomic_api_json.py进行详细分析。这里是原论文地址与github项目地址。

创建工具部分create_tool/Atomic_api_json.py

这段代码的主要目的是使用 tushare 库获取各种金融数据，并将这些数据的描述、示例输入和输出保存到一个 JSON 文件中。以下是对代码的详细解释：

导入必要的库

import tushare as ts
import matplotlib.pyplot as plt
import pandas as pd
import os
import random
from matplotlib.ticker import MaxNLocator
import matplotlib.font_manager as fm
from prettytable import PrettyTable
from blessed import Terminal
import time
from datetime import datetime, timedelta
import numpy as np
import mplfinance as mpf
from prettytable import PrettyTable
from typing import Optional
import matplotlib.font_manager as fm
from matplotlib.lines import Line2D
from typing import Union, Any
import json

这些库提供了数据获取、数据处理、绘图、文件操作等功能。

获取 Tushare API 的 token 并初始化

token = os.getenv("TUSHARE_TOKEN")
pro = ts.pro_api(token)
all_atomic_api = {}

从环境变量中获取 Tushare 的 token，并初始化 pro 对象以便后续调用 API。

定义并获取各种金融数据

代码分多个部分，对于不同类型的金融数据，生成对应的获取数据的api，每个部分的结构大致相同：

1. 获取数据
2. 提取数据的第一行和最后一行
3. 生成列名的中英文对照字典
4. 创建描述数据的字典
5. 将描述数据的字典添加到 all_atomic_api 中

以下是每个部分的详细解释：

财务指标数据

df = pro.fina_indicator(ts_code="600000.SH", start_date="20200104", end_date="20220104", fields=["ts_code", "end_date", "eps", "current_ratio", "quick_ratio", "inv_turn", "netprofit_margin", "grossprofit_margin", "roe", "roa", "roic", "debt_to_assets", "netprofit_yoy", "dt_netprofit_yoy"])
df_sample_str = df.iloc[[0, -1], :].to_string(header=False, index=False)

columns = df.columns.tolist()
columns_means = ['股票代码', '报告期', '每股收益', '流动比率', '速动比率', '存货周转率', '销售净利率', '销售毛利率', '净资产收益率', '总资产净利率', '投入资本回报率', '资产负债率', '净利润同比增长率', '扣非净利润同比增长率']
columns_dict = dict(zip(columns, columns_means))

atomic_api = {"func_desc": None, "Column_name": None, "example_input": None, "output_first_and_last_row": None}
atomic_api["func_desc"] = "财务指标"
atomic_api["Column_name"] = columns_dict
atomic_api["example_input"] = "pro.fina_indicator(ts_code = \"600000.SH\", start_date = \"20200104\", end_date = \"20220104\", fields=[\"ts_code\", \"end_date\", \"eps\", \"current_ratio\", \"quick_ratio\", \"inv_turn\", \"netprofit_margin\", \"grossprofit_margin\", \"roe\", \"roa\", \"roic\", \"debt_to_assets\", \"netprofit_yoy\", \"dt_netprofit_yoy\"])"
atomic_api["output_first_and_last_row"] = df_sample_str

all_atomic_api["pro.fina_indicator"] = atomic_api

这个部分获取了股票代码为 600000.SH 的财务指标数据，并将其描述信息保存到 all_atomic_api 字典中。让我们来看一下最终的结果，其实是生成了一个获取数据的api:

"pro.fina_indicator": {
        "func_desc": "财务指标",
        "Column_name": {
            "ts_code": "股票代码",
            "end_date": "报告期",
            "eps": "每股收益",
            "current_ratio": "流动比率",
            "quick_ratio": "速动比率",
            "inv_turn": "存货周转率",
            "netprofit_margin": "销售净利率",
            "grossprofit_margin": "销售毛利率",
            "roe": "净资产收益率",
            "roa": "总资产净利率",
            "roic": "投入资本回报率",
            "debt_to_assets": "资产负债率",
            "netprofit_yoy": "净利润同比增长率",
            "dt_netprofit_yoy": "扣非净利润同比增长率"
        },
        "example_input": "pro.fina_indicator(ts_code = \"600000.SH\",start_date = \"20200104\",end_date = \"20220104\",fields=[\"ts_code\",\"end_date\",\"eps\",\"current_ratio\",\"quick_ratio\",\"inv_turn\",\"netprofit_margin\",\"grossprofit_margin\",\"roe\",\"roa\",\"roic\",\"debt_to_assets\",\"netprofit_yoy\",\"dt_netprofit_yoy\"])",
        "output_first_and_last_row": "600000.SH 20211231 1.62 None None None 28.1524 None 8.1032 None None 91.6648 -9.1247 -9.26\n600000.SH 20200331 0.56 None None None 31.6289 None 3.0746 None None 92.0170  5.4803  5.23"
    }

获取上市公司基本信息

df = pro.stock_company(ts_code='600230.SH', fields=["ts_code", "exchange", "chairman", "manager", "secretary", "reg_capital", "setup_date", "province", "city", "introduction", "website", "email", "office", "employees", "main_business", "business_scope"])
df_sample_str = df.iloc[0, :].to_string(header=False, index=False)

columns = df.columns.tolist()
columns_means = ['股票代码', '交易所代码', '法人代表', '总经理', '董秘', '注册资本', '注册日期', '所在省份', '所在城市', '公司介绍', '公司主页', '电子邮件', '办公室地址', '员工人数', '主要业务及产品', '经营范围']
columns_dict = dict(zip(columns, columns_means))

atomic_api = {"func_desc": None, "Column_name": None, "example_input": None, "output_first_and_last_row": None}
atomic_api["func_desc"] = "获取上市公司基本信息如公司业务描述,员工人数等基本信息"
atomic_api["Column_name"] = columns_dict
atomic_api["example_input"] = "pro.stock_company(ts_code = '600230.SH', fields=[\"ts_code\", \"exchange\", \"chairman\", \"manager\", \"secretary\", \"reg_capital\", \"setup_date\", \"province\", \"city\", \"introduction\", \"website\", \"email\", \"office\", \"employees\", \"main_business\", \"business_scope\"])"
atomic_api["output_first_and_last_row"] = df_sample_str

all_atomic_api["pro.stock_company"] = atomic_api

这个部分获取了股票代码为 600230.SH 的上市公司基本信息，并将其描述信息保存到 all_atomic_api 字典中。

获取股票每日基本指标

df = pro.daily_basic(ts_code="600230.SH", start_date="20180726", end_date="20200726", fields=["ts_code", "trade_date", "turnover_rate", "turnover_rate_f", "volume_ratio", "pe_ttm", "pb", "ps_ttm", "dv_ttm", "total_share", "float_share", "free_share", "total_mv", "circ_mv"])
df_sample_str = df.iloc[[0, -1], :].to_string(header=False, index=False)

columns = df.columns.tolist()
columns_means = ['股票代码', '交易日期', '换手率(总股本)', '换手率(自由流通股本)', '量比', '市盈率(动态)', '市净率', '市销率(动态)', '股息率(动态)', '总股本', '流通股本', '自由流通股本', '总市值', '流通市值']
columns_dict = dict(zip(columns, columns_means))

atomic_api = {"func_desc": None, "Column_name": None, "example_input": None, "output_first_and_last_row": None}
atomic_api["func_desc"] = "获取股票每日基本指标例如换手率,市盈率市净率股息率等"
atomic_api["Column_name"] = columns_dict
atomic_api["example_input"] = "pro.daily_basic(ts_code = \"600230.SH\", start_date = \"20180726\", end_date = \"20200726\", fields=[\"ts_code\", \"trade_date\", \"turnover_rate\", \"turnover_rate_f\", \"volume_ratio\", \"pe_ttm\", \"pb\", \"ps_ttm\", \"dv_ttm\", \"total_share\", \"float_share\", \"free_share\", \"total_mv\", \"circ_mv\"])"
atomic_api["output_first_and_last_row"] = df_sample_str

all_atomic_api["pro.daily_basic"] = atomic_api

这个部分获取了股票代码为 600230.SH 的每日基本指标数据，并将其描述信息保存到 all_atomic_api 字典中。

获取股票的每日技术指标数据

df = pro.stk_factor(ts_code="600000.SH", start_date="20220520", end_date="20230520", fields=["ts_code", "