crossborder/anhui/gov_commodity_anhui_country.py

from pathlib import Path

import pandas as pd

from anhui import download_dir
from utils import base_country_code, base_mysql
from utils.base_country_code import format_sql_value
from utils.log import log

# 排除地区名单
EXCLUDE_REGIONS = ["亚洲", "非洲", "欧洲", "拉丁美洲", "北美洲", "大洋洲", "南极洲",
                   "东南亚国家联盟", "欧洲联盟", "亚太经济合作组织",
                   "区域全面经济伙伴关系协定（RCEP）成员国", "共建“一带一路”国家和地区"]

def get_df(path):
    file_paths = list(Path(path).glob('*'))
    if not file_paths:
        log.info("未找到任何文件")
        return None
    for file in file_paths:
        if "国别" in file.name:
            log.info(f"处理多文件: {file.name}")
            file_path = Path(path) / file
            return pd.read_excel(file_path, header=None).iloc[6:]
    return None

def process_folder(path):
    year, month = base_country_code.extract_year_month_from_path(path)
    year_month = f'{year}-{month:02d}'

    sql_arr_copy = []
    sql_arr = []
    # try:
    df = get_df(path)
    if df is None:
        log.info("未找到任何文件")
        return None

    country_name_index = 0
    if month == 2:
        col_total_index, col_monthly_export_index, col_monthly_import_index = 1, 3, 5
    else:
        col_total_index, col_monthly_export_index, col_monthly_import_index = 3, 7, 11

    for index, row in df.iterrows():
        # 提取国家名称并去除括号内容
        country_name = str(row.values[country_name_index]).strip()
        if country_name.endswith("）") or country_name.endswith(")"):
            country_name = country_name.rsplit("（")[0] or country_name.rsplit("(")[0]

        # 过滤掉排除地区
        if country_name in EXCLUDE_REGIONS:
            continue

        # 获取国家编码
        country_code = base_country_code.COUNTRY_CODE_MAPPING.get(country_name)
        if not country_code:
            log.info(f"{year_month} 未找到国家 '{country_name}' 对应的编码")
            continue

        # 提取数据并格式化
        monthly_export, monthly_import, monthly_total, yoy_export, yoy_import, yoy_import_export = \
            value_row(row, col_total_index, col_monthly_export_index, col_monthly_import_index)

        if month == 2:
            # 所有总额除2
            year_month_2 = f'{year}-01'
            monthly_import = round(float(monthly_import) / 2, 4)
            monthly_export = round(float(monthly_export) / 2, 4)
            monthly_total = round(float(monthly_total) / 2, 4)
            yoy_import_export, yoy_import, yoy_export = 0, 0, 0
            sql = (f"INSERT INTO t_yujin_crossborder_prov_country_trade "
                   f"(crossborder_year, crossborder_year_month, prov_code, prov_name, country_code, country_name, monthly_total, monthly_export, monthly_import,yoy_import_export, yoy_import, yoy_export, create_time) VALUES "
                   f"('{year}', '{year_month_2}', '340000', '安徽省', '{country_code}', '{country_name}', {format_sql_value(monthly_total)}, {format_sql_value(monthly_export)}, {format_sql_value(monthly_import)}, '{yoy_import_export}', '{yoy_import}', '{yoy_export}', now())"
                   f"ON DUPLICATE KEY UPDATE create_time = now() ;")
            sql_arr_copy.append(sql)

        # 构建 SQL
        sql = (
            f"INSERT INTO t_yujin_crossborder_prov_country_trade "
            f"(crossborder_year, crossborder_year_month, prov_code, prov_name, country_code, country_name, "
            f"monthly_total, monthly_export, monthly_import, yoy_import_export, yoy_import, yoy_export, create_time) "
            f"VALUES ('{year}', '{year_month}', '340000', '安徽省', '{country_code}', '{country_name}', "
            f"{format_sql_value(monthly_total)}, {format_sql_value(monthly_export)}, {format_sql_value(monthly_import)}, '{yoy_import_export}', '{yoy_import}', "
            f"'{yoy_export}', NOW()) ON DUPLICATE KEY UPDATE create_time = now();"
        )
        sql_arr.append(sql)

    log.info(f"√ {year_month} 成功生成 SQL 条数: {len(sql_arr)}")
    # 批量插入数据库
    base_mysql.bulk_insert(sql_arr)
    if month == 2:
        log.info(f"√ {year_month} prov_country_trade 成功生成 SQL 文件 size {len(sql_arr_copy)} ")
        base_mysql.bulk_insert(sql_arr_copy)
    log.info(f"√ {year_month} prov_country_trade SQL 存表完成！\n")




def value_row(row, col_total_index, col_monthly_export_index, col_monthly_import_index):
    def value_special_handler(value):
        if pd.isna(value) or value == "--":
            return float(0)
        else:
            return float(value.strip())

    monthly_total = round(value_special_handler(str(row.values[col_total_index])) * 10000, 4)
    yoy_import_export = 0
    monthly_export = round(value_special_handler(str(row.values[col_monthly_export_index])) * 10000, 4)
    yoy_export = 0
    monthly_import = round(value_special_handler(str(row.values[col_monthly_import_index])) * 10000, 4)
    yoy_import = 0

    return monthly_export, monthly_import, monthly_total, yoy_export, yoy_import, yoy_import_export


def hierarchical_traversal(root_path):
    root = Path(root_path)
    year_dirs = [
        item for item in root.iterdir()
        if item.is_dir() and base_country_code.YEAR_PATTERN.match(item.name)
    ]

    for year_dir in sorted(year_dirs, key=lambda x: x.name, reverse=True):
        log.info(f"\n年份：{year_dir.name} | 省份：anhui")

        month_dirs = []
        for item in year_dir.iterdir():
            if item.is_dir() and base_country_code.MONTH_PATTERN.match(item.name):
                month_dirs.append({"path": item, "month": int(item.name)})

        if month_dirs:
            for md in sorted(month_dirs, key=lambda x: x["month"], reverse=True):
                log.info(f"  月份：{md['month']:02d} | 路径：{md['path']}")
                process_folder(md['path'])


if __name__ == '__main__':
    hierarchical_traversal(download_dir)
    log.info("安徽合肥海关国别所有文件处理完成！")