dqa/loader/file_reader.py - master - nolizard/DQA

dqa/loader/file_reader.py
181 строка5 KB
Mikhail Nyrkov
Initial commit
14 апр 2026, 19:01
14 апр 2026, 19:018911266
"""File Reader module."""

import csv
import os
from pathlib import Path
import sys
import tempfile
from typing import Dict, NoReturn, Optional, Union
import zipfile
from dqa.logger import log_execution

from charset_normalizer import from_bytes
import xmltodict

from .datalab_patch import adapt_path
from .engines.engine_config import ENGINE

from dqa.file_system_utils import ensure_directory_exists

if ENGINE == 'polars':
    from .engines import polars_engine as engine
else:
    from .engines import pandas_engine as engine

# Datatype of returns of the module according to engine.
# Will equal to polars.DataFrame or pandas.Dataframe for example.
ENGINE_TYPE = engine.RETURN_TYPE

NUMBER_OF_ROWS_TO_DEFINE = 50


def define_reading_way(path: Union[str, Path]) -> dict:
    """Offers a way how to process a filepath.
    For the 1st GUI scenario.
    """
    path = Path(adapt_path(path))
    if not path.is_file():
        return {'isfile': False}
    mode = choose_policy(path)
    if mode != 'xlsx':
        return {'isfile': True, 'mode': mode}
    return {'isfile': True, 'mode': mode, 'sheets': __get_xlsx_sheets(path)}


def choose_policy(path: Path) -> Optional[str]:
    """Chooses policy (csv-like or xlsx-like) by file extension."""
    if path.suffix in {'.csv', '.txt'}:
        return 'csv'
    elif path.suffix in {'.xls', '.xlsx'}:
        return 'xlsx'
    return None


def __get_xlsx_sheets(path: Path) -> list:
    """Extracts list of sheets for *.xlsx."""
    with tempfile.TemporaryDirectory(dir='') as dir_to_extract:
        __unzip_xlsx(path, dir_to_extract)

        path_to_xml = dir_to_extract + '/xl/workbook.xml'
        sheets = __get_sheet_names_from_xml(path_to_xml)

    return sheets


def __unzip_xlsx(file_path: Path, dir_to_extract: Path) -> NoReturn:
    """Extracts all from *.xlsx."""
    zip_ref = zipfile.ZipFile(file_path, 'r')
    zip_ref.extractall(dir_to_extract)
    zip_ref.close()


def __get_sheet_names_from_xml(path_to_xml: Path) -> list:
    """Parses sheet names from xml."""
    with open(path_to_xml, 'r', encoding='utf-8') as xml:
        sheets: list = []
        dictionary: Dict[str, Dict[str, Dict]] = xmltodict.parse(xml.read())

        sheets_holder = dictionary['workbook']['sheets']['sheet']
        if isinstance(sheets_holder, dict):
            sheets = [sheets_holder['@name']]
        else:
            for sheet in sheets_holder:
                sheets.append(sheet['@name'])
        return sheets


@log_execution
def read_auto(path: Union[str, Path], n_rows: Optional[int] = None,
              encoding: str = '', sep: str = '',
              sheet_name: Optional[str] = None) -> ENGINE_TYPE:
    """Reads a file by params.
    Main function for 2nd and 3rd GUI scenarios.
    """

    path = Path(path)
    path = adapt_path(path)
    mode = choose_policy(path)
    if mode == 'csv':
        return __read_csv(path, n_rows, encoding, sep)
    elif mode == 'xlsx':
        return __read_excel(path, n_rows, sheet_name)


def __read_csv(path: Path, n_rows: Optional[int],
               encoding: str, sep: str) -> ENGINE_TYPE:
    """Reads csv-like by user params if any,
    otherwise determine by itself.
    """
    if not encoding or encoding == 'auto':
        encoding = __define_encoding(path)
    if not sep or sep == 'auto':
        sep = __define_separator(path, encoding)
    return engine.read_csv(path, sep, encoding, n_rows)


def __read_excel(path: Path, n_rows: Optional[int] = None,
                 sheet_name: Optional[str] = None) -> ENGINE_TYPE:
    """Reads xlsx-like from user sheet, otherwise from 1st."""
    data = engine.read_excel(path, sheet_name)
    if n_rows:
        return data.head(n_rows)
    return data


def __define_encoding(path: Path) -> str:
    """Defines an encoding of csv-like using charset_normalizer lib."""
    with path.open('rb') as file:
        lines = [file.readline() for _ in range(NUMBER_OF_ROWS_TO_DEFINE)]
        raw = b''.join(line for line in lines if line.strip())
    return from_bytes(raw).best().encoding


def __define_separator(path: Path, encoding: str) -> str:
    """Defines a separator of csv-like using csv lib."""
    with path.open(mode='r', encoding=encoding) as file:
        data = ''
        number_of_completed_lines = 0
        while number_of_completed_lines <= NUMBER_OF_ROWS_TO_DEFINE:
            line = file.readline()
            if line.strip():
                data += line
            number_of_completed_lines += 1
        dialect = csv.Sniffer().sniff(data)
        return str(dialect.delimiter)


def save_df_csv(df: ENGINE_TYPE, path: Path):
    df.to_csv(path, index=False)


def save_df_excel(df: ENGINE_TYPE, path: Path):
    df.to_excel(path, index=False)


def save_file(content: str, path: Path):
    with open(path, 'w', encoding='utf-8') as file:
        file.write(content)


def dump_input(path: str):
    """Save log path into `inputs`."""
    try:
        ensure_directory_exists('inputs')

        with open(os.path.join('inputs', 'log_path.txt'), 'w') as txt_file:
            txt_file.write(path)

    except OSError as e:
        print(f"Error dumping: {e}", file=sys.stderr)


def load_input() -> Union[str, bool]:
    """Read log path from `inputs`."""
    try:
        with open(os.path.join('inputs', 'log_path.txt'), 'r') as txt_file:
                path = txt_file.read()
        return read_auto(path)

    except OSError as e:
        print(f"Error dumping: {e}", file=sys.stderr)
        return False