easy-ml/data/loader/data_loader_excel.py

# -*- coding: utf-8 -*-
"""
@author: yq
@time: 2024/10/31
@desc: 
"""
import pandas as pd

from commom import get_logger
from .data_loader_base import DataLoaderBase

logger = get_logger()


class DataLoaderExcel(DataLoaderBase):
    def __init__(self, ):
        pass

    def get_connect(self):
        pass

    def close_connect(self):
        pass

    def get_data(self, file_path: str, sheet_name: str = 0) -> pd.DataFrame:
        df: pd.DataFrame = pd.DataFrame()
        if ".xlsx" in file_path:
            df = pd.read_excel(file_path, sheet_name=sheet_name, index_col=False)
        elif ".csv" in file_path:
            df = pd.read_csv(file_path)
        columns = df.columns.to_list()
        columns_new = []
        for idx, column in enumerate(columns):
            column = str(column)
            if idx != 0 and "Unnamed:" in column:
                columns_new.append(columns_new[-1])
            else:
                columns_new.append(column)
        df.columns = columns_new
        return df