energy-data-trans/utils/trans_methods.py

# -*- coding: utf-8 -*-
# @Time    : 2024/5/16
# @Author  : 魏志亮
import os
import re
import warnings

import chardet
import pandas as pd

from utils.log.trans_log import trans_print

warnings.filterwarnings("ignore")


# 获取文件编码
def detect_file_encoding(filename):
    # 读取文件的前1000个字节（足够用于大多数编码检测）
    with open(filename, 'rb') as f:
        rawdata = f.read(1000)
    result = chardet.detect(rawdata)
    return result['encoding']


# 读取数据到df
def read_file_to_df(file_path, read_cols=list()):
    trans_print('开始读取文件', file_path)
    df = pd.DataFrame()
    encoding = detect_file_encoding(file_path)

    if str(file_path).lower().endswith("csv"):
        if read_cols:
            df = pd.read_csv(file_path, encoding=encoding, usecols=read_cols)
        else:
            df = pd.read_csv(file_path, encoding=encoding)
    else:
        xls = pd.ExcelFile(file_path)
        # 获取所有的sheet名称
        sheet_names = xls.sheet_names
        for sheet in sheet_names:
            if read_cols:
                df = pd.concat([df, pd.read_excel(xls, sheet_name=sheet, usecols=read_cols)])
            else:
                df = pd.concat([df, pd.read_excel(xls, sheet_name=sheet)])

    trans_print('文件读取成功', file_path, '文件数量', df.shape)

    return df


def __build_directory_dict(directory_dict, path):
    # 遍历目录下的所有项
    for item in os.listdir(path):
        item_path = os.path.join(path, item)
        if os.path.isdir(item_path):
            __build_directory_dict(directory_dict, item_path)
        elif os.path.isfile(item_path):
            if path not in directory_dict:
                directory_dict[path] = []

            types = ['xls', 'xlsx', 'csv']
            if str(item_path).split(".")[-1] in types:
                if str(item_path).count("~$") == 0:
                    directory_dict[path].append(item_path)

    # 读取所有文件


# 读取路径下所有的excel文件
def read_excel_files(read_path):
    directory_dict = {}
    __build_directory_dict(directory_dict, read_path)

    return [path for paths in directory_dict.values() for path in paths if path]


# 创建路径
def create_file_path(path, is_file_path=False):
    if is_file_path:
        path = os.path.dirname(path)

    if not os.path.exists(path):
        os.makedirs(path)


# 格式化风机名称
def generate_turbine_name(turbine_name='F0001', prefix='F'):
    strinfo = re.compile(r"[\D*]")
    name = strinfo.sub('', str(turbine_name))
    return prefix + str(int(name)).zfill(3)