Repository: topoteretes/cognee
File: cognee/infrastructure/files/utils/guess_file_type.py (Python)

Code before:

    Parameters:
    -----------

        - file (BinaryIO): A binary file stream to analyze for determining the file type.

    Returns:
    --------

        - filetype.Type: The guessed file type, represented as filetype.Type.
    """
    file_type = filetype.guess(file)

    # If file type could not be determined consider it a plain text file as they don't have magic number encoding
    if file_type is None:
        from filetype.types.base import Type

        file_type = Type("text/plain", "txt")

    if file_type is None:
        raise FileTypeException(f"Unknown file detected: {file.name}.")

    return file_type


class CsvFileType(filetype.Type):
    """
    Match CSV file types based on MIME type and extension.

    Public methods:
    - match

    Instance variables:
    - MIME: The MIME type of the CSV.
    - EXTENSION: The file extension of the CSV.
    """

    MIME = "text/csv"
    EXTENSION = "csv"

    def __init__(self):
        super().__init__(mime=self.MIME, extension=self.EXTENSION)

    def match(self, buf):
        """
        Determine if the given buffer contains csv content.

        Parameters:
        -----------

            - buf: The buffer to check for csv content.

Review comment (suggestion):
We use custom loaders now, we don't expand this file anymore.
You add a custom loader to cognee by using `use_loader` hook.

Code after:
import io
from pathlib import Path
from typing import BinaryIO, Optional, Any
import filetype
from tempfile import SpooledTemporaryFile
from filetype.types.base import Type


class FileTypeException(Exception):
    """
    Represents an exception for file type errors.

    This exception is raised when an invalid file type is encountered. It includes a custom
    message to describe the error.

    Parameters:
    -----------

        - message (str): The message describing the exception error.
    """

    message: str

    def __init__(self, message: str):
        self.message = message


def guess_file_type(file: BinaryIO, name: Optional[str] = None) -> filetype.Type:
    """
    Guess the file type from the given binary file stream.

    If the file type cannot be determined from content, attempts to infer from extension.
    If still unable to determine, raise a FileTypeException with an appropriate message.

    Parameters:
    -----------

        - file (BinaryIO): A binary file stream to analyze for determining the file type.

    Returns:
    --------

        - filetype.Type: The guessed file type, represented as filetype.Type.
    """

    # Note: If file has .txt or .text extension, consider it a plain text file as filetype.guess may not detect it properly
    # as it contains no magic number encoding
    ext = None
    if isinstance(file, str):
        ext = Path(file).suffix
    elif name is not None:
        ext = Path(name).suffix

    if ext in [".txt", ".text"]:
        file_type = Type("text/plain", "txt")
        return file_type

    if ext in [".csv"]:
        file_type = Type("text/csv", "csv")
        return file_type

    file_type = filetype.guess(file)

    # If file type could not be determined consider it a plain text file as they don't have magic number encoding
    if file_type is None:
        file_type = Type("text/plain", "txt")

    if file_type is None:
        raise FileTypeException(f"Unknown file detected: {file.name}.")

    return file_type