unstructured-io--unstructured/unstructured/file_utils/encoding.py

from typing import IO, Optional, Tuple, Union

from charset_normalizer import detect

from unstructured.errors import UnprocessableEntityError
from unstructured.partition.common.common import convert_to_bytes

ENCODE_REC_THRESHOLD = 0.8

# popular encodings from https://en.wikipedia.org/wiki/Popularity_of_text_encodings
COMMON_ENCODINGS = [
    "utf_8",
    "iso_8859_1",
    "iso_8859_6",
    "iso_8859_8",
    "ascii",
    "big5",
    "utf_16",
    "utf_16_be",
    "utf_16_le",
    "utf_32",
    "utf_32_be",
    "utf_32_le",
    "euc_jis_2004",
    "euc_jisx0213",
    "euc_jp",
    "euc_kr",
    "gb18030",
    "shift_jis",
    "shift_jis_2004",
    "shift_jisx0213",
]


def format_encoding_str(encoding: str) -> str:
    """Format input encoding string (e.g., `utf-8`, `iso-8859-1`, etc).
    Parameters
    ----------
    encoding
        The encoding string to be formatted (e.g., `UTF-8`, `utf_8`, `ISO-8859-1`, `iso_8859_1`,
        etc).
    """
    formatted_encoding = encoding.lower().replace("_", "-")

    # Special case for Arabic and Hebrew charsets with directional annotations
    annotated_encodings = ["iso-8859-6-i", "iso-8859-6-e", "iso-8859-8-i", "iso-8859-8-e"]
    if formatted_encoding in annotated_encodings:
        formatted_encoding = formatted_encoding[:-2]  # remove the annotation

    return formatted_encoding


def validate_encoding(encoding: str) -> bool:
    """Checks if an encoding string is valid. Helps to avoid errors in cases where
    invalid encodings are extracted from malformed documents."""
    for common_encoding in COMMON_ENCODINGS:
        if format_encoding_str(common_encoding) == format_encoding_str(encoding):
            return True
    return False


def detect_file_encoding(
    filename: str = "",
    file: Optional[Union[bytes, IO[bytes]]] = None,
) -> Tuple[str, str]:
    if filename:
        with open(filename, "rb") as f:
            byte_data = f.read()
    elif file:
        byte_data = convert_to_bytes(file)
    else:
        raise FileNotFoundError("No filename nor file were specified")

    result = detect(byte_data)
    encoding = result["encoding"]
    confidence = result["confidence"]

    if encoding is None or confidence is None or confidence < ENCODE_REC_THRESHOLD:
        # Encoding detection failed, fallback to predefined encodings
        for enc in COMMON_ENCODINGS:
            try:
                if filename:
                    with open(filename, encoding=enc) as f:
                        file_text = f.read()
                else:
                    file_text = byte_data.decode(enc)
                encoding = enc
                break
            except (UnicodeDecodeError, UnicodeError):
                continue
        else:
            # NOTE: Use UnprocessableEntityError instead of UnicodeDecodeError to avoid
            # logging the entire file content. UnicodeDecodeError automatically stores
            # the complete input data, which can be problematic for large files.
            raise UnprocessableEntityError(
                "Unable to determine file encoding after trying all common encodings. "
                "File may be corrupted or in an unsupported format."
            ) from None

    else:
        # NOTE: Catch UnicodeDecodeError to avoid logging the entire file content.
        # UnicodeDecodeError automatically stores the complete input data in its
        # 'object' attribute, which can cause issues with large files in logging
        # and error reporting systems.
        try:
            file_text = byte_data.decode(encoding)
        except (UnicodeDecodeError, UnicodeError):
            raise UnprocessableEntityError(
                f"File encoding detection failed: detected '{encoding}' but decode failed. "
                f"File may be corrupted or in an unsupported format."
            ) from None

    formatted_encoding = format_encoding_str(encoding)

    return formatted_encoding, file_text


def read_txt_file(
    filename: str = "",
    file: Optional[Union[bytes, IO[bytes]]] = None,
    encoding: Optional[str] = None,
) -> Tuple[str, str]:
    """Extracts document metadata from a plain text document."""
    if filename:
        if encoding:
            formatted_encoding = format_encoding_str(encoding)
            with open(filename, encoding=formatted_encoding) as f:
                try:
                    file_text = f.read()
                except (UnicodeDecodeError, UnicodeError) as error:
                    raise error
        else:
            formatted_encoding, file_text = detect_file_encoding(filename)
    elif file:
        if encoding:
            formatted_encoding = format_encoding_str(encoding)
            try:
                file_content = file if isinstance(file, bytes) else file.read()
                if isinstance(file_content, bytes):
                    file_text = file_content.decode(formatted_encoding)
                else:
                    file_text = file_content
            except (UnicodeDecodeError, UnicodeError) as error:
                raise error
        else:
            formatted_encoding, file_text = detect_file_encoding(file=file)
    else:
        raise FileNotFoundError("No filename was specified")

    return formatted_encoding, file_text