办公自动化_帮人事做个表格拆分

# coding=utf-8
import pandas as pd
import openpyxl as op
import os
from shutil import rmtree


def groupby_data(data, chaifen, file_a):
    new_data = data.groupby(chaifen)
    print("新文件生成会在程序所在目录下")
    for i in new_data:
        wb_name = '+'.join([str(x) for x in i[0]])
        wb_name = wb_name.replace("/", "&")
        save_excel(wb_name=wb_name, data=i[1], file=file_a)


def save_excel(wb_name, data, file):
    excel_file = file + "\\" + wb_name + '.xlsx'
    data.to_excel(excel_file,
                  encoding='utf-8',
                  startrow=0,
                  startcol=0,
                  index=False)
    print("文件[%s]已生成" % excel_file)


def get_excel(file, header=0, sheet_name=0, skip_head=None, skip_foot=0):
    data = pd.read_excel(file,
                         sheet_name=sheet_name,
                         header=header,
                         skiprows=skip_head,
                         skipfooter=skip_foot)
    return data


def select_sheet_name(file):
    wb = op.load_workbook(file)
    sheet_name = wb.sheetnames
    print("如果表名大于一个，则说明存在隐藏表，故输入相应表名即可\n", sheet_name)


if __name__ == "__main__":
    try:
        a = r"C:\Users\78563\Desktop\拆分结果"
        if os.path.isdir(a) is False:
            os.mkdir(a)
        file = input("请输入文件地址:")
        select_sheet_name(file)
        sheet_name = input("请选择表:")
        skip_head = int(input("跳过头几行:"))
        skip_foot = int(input("跳过尾几行:"))
        data = get_excel(file=file,
                         header=None,
                         sheet_name=sheet_name,
                         skip_head=skip_head,
                         skip_foot=skip_foot)
        excel_file = r'D:\备份.xlsx'
        data.to_excel(excel_file,
                      encoding='utf-8',
                      startrow=0,
                      startcol=0,
                      index=False)

        data = get_excel(excel_file, header=0, sheet_name=0)
        chaifen = []
        print("输入666则开始筛选")
        while True:
            col_name = int(input("请输入要进行拆分的列序号(默认从1开始):"))
            col_name -= 1
            if col_name == 665:
                break
            chaifen.append(col_name)
            print(chaifen)
        groupby_data(data=data, chaifen=chaifen, file_a=a)
    except FileNotFoundError:
        print("文件存储错误")
        rmtree(a)
    except OSError:
        print("文件存储错误")
        rmtree(a)
    else:
        print("程序执行完毕")
    finally:
        os.remove(excel_file)