lfx-nightly 0.1.11.dev0__py3-none-any.whl

lfx/base/data/base_file.py

@@ -0,0 +1,685 @@
+import ast
+import json
+import shutil
+import tarfile
+from abc import ABC, abstractmethod
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import TYPE_CHECKING
+from zipfile import ZipFile, is_zipfile
+
+import pandas as pd
+
+from lfx.custom.custom_component.component import Component
+from lfx.io import BoolInput, FileInput, HandleInput, Output, StrInput
+from lfx.schema.data import Data
+from lfx.schema.dataframe import DataFrame
+from lfx.schema.message import Message
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+class BaseFileComponent(Component, ABC):
+    """Base class for handling file processing components.
+
+    This class provides common functionality for resolving, validating, and
+    processing file paths. Child classes must define valid file extensions
+    and implement the `process_files` method.
+    """
+
+    class BaseFile:
+        """Internal class to represent a file with additional metadata."""
+
+        def __init__(
+            self,
+            data: Data | list[Data],
+            path: Path,
+            *,
+            delete_after_processing: bool = False,
+            silent_errors: bool = False,
+        ):
+            self._data = data if isinstance(data, list) else [data]
+            self.path = path
+            self.delete_after_processing = delete_after_processing
+            self._silent_errors = silent_errors
+
+        @property
+        def data(self) -> list[Data]:
+            return self._data or []
+
+        @data.setter
+        def data(self, value: Data | list[Data]):
+            if isinstance(value, Data):
+                self._data = [value]
+            elif isinstance(value, list) and all(isinstance(item, Data) for item in value):
+                self._data = value
+            else:
+                msg = f"data must be a Data object or a list of Data objects. Got: {type(value)}"
+                if not self._silent_errors:
+                    raise ValueError(msg)
+
+        def merge_data(self, new_data: Data | list[Data] | None) -> list[Data]:
+            r"""Generate a new list of Data objects by merging `new_data` into the current `data`.
+
+            Args:
+                new_data (Data | list[Data] | None): The new Data object(s) to merge into each existing Data object.
+                    If None, the current `data` is returned unchanged.
+
+            Returns:
+                list[Data]: A new list of Data objects with `new_data` merged.
+            """
+            if new_data is None:
+                return self.data
+
+            if isinstance(new_data, Data):
+                new_data_list = [new_data]
+            elif isinstance(new_data, list) and all(isinstance(item, Data) for item in new_data):
+                new_data_list = new_data
+            else:
+                msg = "new_data must be a Data object, a list of Data objects, or None."
+                if not self._silent_errors:
+                    raise ValueError(msg)
+                return self.data
+
+            return [
+                Data(data={**data.data, **new_data_item.data}) for data in self.data for new_data_item in new_data_list
+            ]
+
+        def __str__(self):
+            if len(self.data) == 0:
+                text_preview = ""
+            elif len(self.data) == 1:
+                max_text_length = 50
+                text_preview = self.data.get_text()[:max_text_length]
+                if len(self.data.get_text()) > max_text_length:
+                    text_preview += "..."
+                text_preview = f"text_preview='{text_preview}'"
+            else:
+                text_preview = f"{len(self.data)} data objects"
+            return f"BaseFile(path={self.path}, delete_after_processing={self.delete_after_processing}, {text_preview}"
+
+    # Subclasses can override these class variables
+    VALID_EXTENSIONS: list[str] = []  # To be overridden by child classes
+    IGNORE_STARTS_WITH = [".", "__MACOSX"]
+
+    SERVER_FILE_PATH_FIELDNAME = "file_path"
+    SUPPORTED_BUNDLE_EXTENSIONS = ["zip", "tar", "tgz", "bz2", "gz"]
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # Dynamically update FileInput to include valid extensions and bundles
+        self.get_base_inputs()[0].file_types = [
+            *self.valid_extensions,
+            *self.SUPPORTED_BUNDLE_EXTENSIONS,
+        ]
+
+        file_types = ", ".join(self.valid_extensions)
+        bundles = ", ".join(self.SUPPORTED_BUNDLE_EXTENSIONS)
+        self.get_base_inputs()[
+            0
+        ].info = f"Supported file extensions: {file_types}; optionally bundled in file extensions: {bundles}"
+
+    _base_inputs = [
+        FileInput(
+            name="path",
+            display_name="Files",
+            fileTypes=[],  # Dynamically set in __init__
+            info="",  # Dynamically set in __init__
+            required=False,
+            list=True,
+            value=[],
+        ),
+        HandleInput(
+            name="file_path",
+            display_name="Server File Path",
+            info=(
+                f"Data object with a '{SERVER_FILE_PATH_FIELDNAME}' property pointing to server file"
+                " or a Message object with a path to the file. Supercedes 'Path' but supports same file types."
+            ),
+            required=False,
+            input_types=["Data", "Message"],
+            is_list=True,
+            advanced=True,
+        ),
+        StrInput(
+            name="separator",
+            display_name="Separator",
+            value="\n\n",
+            show=True,
+            info="Specify the separator to use between multiple outputs in Message format.",
+            advanced=True,
+        ),
+        BoolInput(
+            name="silent_errors",
+            display_name="Silent Errors",
+            advanced=True,
+            info="If true, errors will not raise an exception.",
+        ),
+        BoolInput(
+            name="delete_server_file_after_processing",
+            display_name="Delete Server File After Processing",
+            advanced=True,
+            value=True,
+            info="If true, the Server File Path will be deleted after processing.",
+        ),
+        BoolInput(
+            name="ignore_unsupported_extensions",
+            display_name="Ignore Unsupported Extensions",
+            advanced=True,
+            value=True,
+            info="If true, files with unsupported extensions will not be processed.",
+        ),
+        BoolInput(
+            name="ignore_unspecified_files",
+            display_name="Ignore Unspecified Files",
+            advanced=True,
+            value=False,
+            info=f"If true, Data with no '{SERVER_FILE_PATH_FIELDNAME}' property will be ignored.",
+        ),
+    ]
+
+    _base_outputs = [
+        Output(display_name="Files", name="dataframe", method="load_files"),
+    ]
+
+    @abstractmethod
+    def process_files(self, file_list: list[BaseFile]) -> list[BaseFile]:
+        """Processes a list of files.
+
+        Args:
+            file_list (list[BaseFile]): A list of file objects.
+
+        Returns:
+            list[BaseFile]: A list of BaseFile objects with updated `data`.
+        """
+
+    def load_files_base(self) -> list[Data]:
+        """Loads and parses file(s), including unpacked file bundles.
+
+        Returns:
+            list[Data]: Parsed data from the processed files.
+        """
+        self._temp_dirs: list[TemporaryDirectory] = []
+        final_files = []  # Initialize to avoid UnboundLocalError
+        try:
+            # Step 1: Validate the provided paths
+            files = self._validate_and_resolve_paths()
+
+            # Step 2: Handle bundles recursively
+            all_files = self._unpack_and_collect_files(files)
+
+            # Step 3: Final validation of file types
+            final_files = self._filter_and_mark_files(all_files)
+
+            # Step 4: Process files
+            processed_files = self.process_files(final_files)
+
+            # Extract and flatten Data objects to return
+            return [data for file in processed_files for data in file.data if file.data]
+
+        finally:
+            # Delete temporary directories
+            for temp_dir in self._temp_dirs:
+                temp_dir.cleanup()
+
+            # Delete files marked for deletion
+            for file in final_files:
+                if file.delete_after_processing and file.path.exists():
+                    if file.path.is_dir():
+                        shutil.rmtree(file.path)
+                    else:
+                        file.path.unlink()
+
+    def load_files_core(self) -> list[Data]:
+        """Load files and return as Data objects.
+
+        Returns:
+            list[Data]: List of Data objects from all files
+        """
+        data_list = self.load_files_base()
+        if not data_list:
+            return [Data()]
+        return data_list
+
+    def load_files_message(self) -> Message:
+        """Load files and return as Message.
+
+        Returns:
+            Message: Message containing all file data
+        """
+        data_list = self.load_files_core()
+        if not data_list:
+            return Message()  # No data -> empty message
+
+        sep: str = getattr(self, "separator", "\n\n") or "\n\n"
+
+        parts: list[str] = []
+        for d in data_list:
+            # Prefer explicit text if available, fall back to full dict, lastly str()
+            text = (getattr(d, "get_text", lambda: None)() or d.data.get("text")) if isinstance(d.data, dict) else None
+            parts.append(text if text is not None else str(d))
+
+        return Message(text=sep.join(parts))
+
+    def load_files_path(self) -> Message:
+        """Returns a Message containing file paths from loaded files.
+
+        Returns:
+            Message: Message containing file paths
+        """
+        files = self._validate_and_resolve_paths()
+        paths = [file.path.as_posix() for file in files if file.path.exists()]
+
+        return Message(text="\n".join(paths) if paths else "")
+
+    def load_files_structured_helper(self, file_path: str) -> list[dict] | None:
+        if not file_path:
+            return None
+
+        # Map file extensions to pandas read functions with type annotation
+        file_readers: dict[str, Callable[[str], pd.DataFrame]] = {
+            ".csv": pd.read_csv,
+            ".xlsx": pd.read_excel,
+            ".parquet": pd.read_parquet,
+            # TODO: sqlite and json support?
+        }
+
+        # Get file extension in lowercase
+        ext = Path(file_path).suffix.lower()
+
+        # Get the appropriate reader function or None
+        reader = file_readers.get(ext)
+
+        if reader:
+            result = reader(file_path)  # MyPy now knows reader is callable
+            return result.to_dict("records")
+
+        return None
+
+    def load_files_structured(self) -> DataFrame:
+        """Load files and return as DataFrame with structured content.
+
+        Returns:
+            DataFrame: DataFrame containing structured content from all files
+        """
+        data_list = self.load_files_core()
+        if not data_list:
+            return DataFrame()
+
+        # Get the file path from the first Data object
+        file_path = data_list[0].data.get(self.SERVER_FILE_PATH_FIELDNAME, None)
+
+        # If file_path is provided and is a CSV, read it directly
+        if file_path and str(file_path).lower().endswith((".csv", ".xlsx", ".parquet")):
+            rows = self.load_files_structured_helper(file_path)
+        else:
+            # Convert Data objects to a list of dictionaries
+            # TODO: Parse according to docling standards
+            rows = [data_list[0].data]
+
+        self.status = DataFrame(rows)
+
+        return DataFrame(rows)
+
+    def parse_string_to_dict(self, s: str) -> dict:
+        # Try JSON first (handles true/false/null)
+        try:
+            result = json.loads(s)
+            if isinstance(result, dict):
+                return result
+        except json.JSONDecodeError:
+            pass
+
+        # Fall back to Python literal evaluation
+        try:
+            result = ast.literal_eval(s)
+            if isinstance(result, dict):
+                return result
+        except (SyntaxError, ValueError):
+            pass
+
+        # If all parsing fails, return the fallback
+        return {"value": s}
+
+    def load_files_json(self) -> Data:
+        """Load files and return as a single Data object containing JSON content.
+
+        Returns:
+            Data: Data object containing JSON content from all files
+        """
+        data_list = self.load_files_core()
+        if not data_list:
+            return Data()
+
+        # Grab the JSON data
+        json_data = data_list[0].data[data_list[0].text_key]
+        json_data = self.parse_string_to_dict(json_data)
+
+        self.status = Data(data=json_data)
+
+        return Data(data=json_data)
+
+    def load_files(self) -> DataFrame:
+        """Load files and return as DataFrame.
+
+        Returns:
+            DataFrame: DataFrame containing all file data
+        """
+        data_list = self.load_files_core()
+        if not data_list:
+            return DataFrame()
+
+        # Convert Data objects to a list of dictionaries
+        all_rows = []
+        for data in data_list:
+            file_path = data.data.get(self.SERVER_FILE_PATH_FIELDNAME)
+            row = dict(data.data) if data.data else {}
+
+            # Add text if available, otherwise use the data's text property
+            if "text" in data.data:
+                row["text"] = data.data["text"]
+            if file_path:
+                row["file_path"] = file_path
+            all_rows.append(row)
+
+        self.status = DataFrame(all_rows)
+
+        return DataFrame(all_rows)
+
+    @property
+    def valid_extensions(self) -> list[str]:
+        """Returns valid file extensions for the class.
+
+        This property can be overridden by child classes to provide specific
+        extensions.
+
+        Returns:
+            list[str]: A list of valid file extensions without the leading dot.
+        """
+        return self.VALID_EXTENSIONS
+
+    @property
+    def ignore_starts_with(self) -> list[str]:
+        """Returns prefixes to ignore when unpacking file bundles.
+
+        Returns:
+            list[str]: A list of prefixes to ignore when unpacking file bundles.
+        """
+        return self.IGNORE_STARTS_WITH
+
+    def rollup_data(
+        self,
+        base_files: list[BaseFile],
+        data_list: list[Data | None],
+        path_field: str = SERVER_FILE_PATH_FIELDNAME,
+    ) -> list[BaseFile]:
+        r"""Rolls up Data objects into corresponding BaseFile objects in order given by `base_files`.
+
+        Args:
+            base_files (list[BaseFile]): The original BaseFile objects.
+            data_list (list[Data | None]): The list of data to be aggregated into the BaseFile objects.
+            path_field (str): The field name on the data_list objects that holds the file path as a string.
+
+        Returns:
+            list[BaseFile]: A new list of BaseFile objects with merged `data` attributes.
+        """
+
+        def _build_data_dict(data_list: list[Data | None], data_list_field: str) -> dict[str, list[Data]]:
+            """Builds a dictionary grouping Data objects by a specified field."""
+            data_dict: dict[str, list[Data]] = {}
+            for data in data_list:
+                if data is None:
+                    continue
+                key = data.data.get(data_list_field)
+                if key is None:
+                    msg = f"Data object missing required field '{data_list_field}': {data}"
+                    self.log(msg)
+                    if not self.silent_errors:
+                        msg = f"Data object missing required field '{data_list_field}': {data}"
+                        self.log(msg)
+                        raise ValueError(msg)
+                    continue
+                data_dict.setdefault(key, []).append(data)
+            return data_dict
+
+        # Build the data dictionary from the provided data_list
+        data_dict = _build_data_dict(data_list, path_field)
+
+        # Generate the updated list of BaseFile objects, preserving the order of base_files
+        updated_base_files = []
+        for base_file in base_files:
+            new_data_list = data_dict.get(str(base_file.path), [])
+            merged_data_list = base_file.merge_data(new_data_list)
+            updated_base_files.append(
+                BaseFileComponent.BaseFile(
+                    data=merged_data_list,
+                    path=base_file.path,
+                    delete_after_processing=base_file.delete_after_processing,
+                )
+            )
+
+        return updated_base_files
+
+    def _file_path_as_list(self) -> list[Data]:
+        file_path = self.file_path
+        if not file_path:
+            return []
+
+        def _message_to_data(message: Message) -> Data:
+            return Data(**{self.SERVER_FILE_PATH_FIELDNAME: message.text})
+
+        if isinstance(file_path, Data):
+            file_path = [file_path]
+        elif isinstance(file_path, Message):
+            file_path = [_message_to_data(file_path)]
+        elif not isinstance(file_path, list):
+            msg = f"Expected list of Data objects in file_path but got {type(file_path)}."
+            self.log(msg)
+            if not self.silent_errors:
+                raise ValueError(msg)
+            return []
+
+        file_paths = []
+        for obj in file_path:
+            data_obj = _message_to_data(obj) if isinstance(obj, Message) else obj
+
+            if not isinstance(data_obj, Data):
+                msg = f"Expected Data object in file_path but got {type(data_obj)}."
+                self.log(msg)
+                if not self.silent_errors:
+                    raise ValueError(msg)
+                continue
+            file_paths.append(data_obj)
+
+        return file_paths
+
+    def _validate_and_resolve_paths(self) -> list[BaseFile]:
+        """Validate that all input paths exist and are valid, and create BaseFile instances.
+
+        Returns:
+            list[BaseFile]: A list of valid BaseFile instances.
+
+        Raises:
+            ValueError: If any path does not exist.
+        """
+        resolved_files = []
+
+        def add_file(data: Data, path: str | Path, *, delete_after_processing: bool):
+            resolved_path = Path(self.resolve_path(str(path)))
+
+            if not resolved_path.exists():
+                msg = f"File or directory not found: {path}"
+                self.log(msg)
+                if not self.silent_errors:
+                    raise ValueError(msg)
+            resolved_files.append(
+                BaseFileComponent.BaseFile(data, resolved_path, delete_after_processing=delete_after_processing)
+            )
+
+        file_path = self._file_path_as_list()
+
+        if self.path and not file_path:
+            # Wrap self.path into a Data object
+            if isinstance(self.path, list):
+                for path in self.path:
+                    data_obj = Data(data={self.SERVER_FILE_PATH_FIELDNAME: path})
+                    add_file(data=data_obj, path=path, delete_after_processing=False)
+            else:
+                data_obj = Data(data={self.SERVER_FILE_PATH_FIELDNAME: self.path})
+                add_file(data=data_obj, path=self.path, delete_after_processing=False)
+        elif file_path:
+            for obj in file_path:
+                server_file_path = obj.data.get(self.SERVER_FILE_PATH_FIELDNAME)
+                if server_file_path:
+                    add_file(
+                        data=obj,
+                        path=server_file_path,
+                        delete_after_processing=self.delete_server_file_after_processing,
+                    )
+                elif not self.ignore_unspecified_files:
+                    msg = f"Data object missing '{self.SERVER_FILE_PATH_FIELDNAME}' property."
+                    self.log(msg)
+                    if not self.silent_errors:
+                        raise ValueError(msg)
+                else:
+                    msg = f"Ignoring Data object missing '{self.SERVER_FILE_PATH_FIELDNAME}' property:\n{obj}"
+                    self.log(msg)
+
+        return resolved_files
+
+    def _unpack_and_collect_files(self, files: list[BaseFile]) -> list[BaseFile]:
+        """Recursively unpack bundles and collect files into BaseFile instances.
+
+        Args:
+            files (list[BaseFile]): List of BaseFile instances to process.
+
+        Returns:
+            list[BaseFile]: Updated list of BaseFile instances.
+        """
+        collected_files = []
+
+        for file in files:
+            path = file.path
+            delete_after_processing = file.delete_after_processing
+            data = file.data
+
+            if path.is_dir():
+                # Recurse into directories
+                collected_files.extend(
+                    [
+                        BaseFileComponent.BaseFile(
+                            data,
+                            sub_path,
+                            delete_after_processing=delete_after_processing,
+                        )
+                        for sub_path in path.rglob("*")
+                        if sub_path.is_file()
+                    ]
+                )
+            elif path.suffix[1:] in self.SUPPORTED_BUNDLE_EXTENSIONS:
+                # Unpack supported bundles
+                temp_dir = TemporaryDirectory()
+                self._temp_dirs.append(temp_dir)
+                temp_dir_path = Path(temp_dir.name)
+                self._unpack_bundle(path, temp_dir_path)
+                subpaths = list(temp_dir_path.iterdir())
+                self.log(f"Unpacked bundle {path.name} into {subpaths}")
+                collected_files.extend(
+                    [
+                        BaseFileComponent.BaseFile(
+                            data,
+                            sub_path,
+                            delete_after_processing=delete_after_processing,
+                        )
+                        for sub_path in subpaths
+                    ]
+                )
+            else:
+                collected_files.append(file)
+
+        # Recurse again if any directories or bundles are left in the list
+        if any(
+            file.path.is_dir() or file.path.suffix[1:] in self.SUPPORTED_BUNDLE_EXTENSIONS for file in collected_files
+        ):
+            return self._unpack_and_collect_files(collected_files)
+
+        return collected_files
+
+    def _unpack_bundle(self, bundle_path: Path, output_dir: Path):
+        """Unpack a bundle into a temporary directory.
+
+        Args:
+            bundle_path (Path): Path to the bundle.
+            output_dir (Path): Directory where files will be extracted.
+
+        Raises:
+            ValueError: If the bundle format is unsupported or cannot be read.
+        """
+
+        def _safe_extract_zip(bundle: ZipFile, output_dir: Path):
+            """Safely extract ZIP files."""
+            for member in bundle.namelist():
+                member_path = output_dir / member
+                # Ensure no path traversal outside `output_dir`
+                if not member_path.resolve().is_relative_to(output_dir.resolve()):
+                    msg = f"Attempted Path Traversal in ZIP File: {member}"
+                    raise ValueError(msg)
+                bundle.extract(member, path=output_dir)
+
+        def _safe_extract_tar(bundle: tarfile.TarFile, output_dir: Path):
+            """Safely extract TAR files."""
+            for member in bundle.getmembers():
+                member_path = output_dir / member.name
+                # Ensure no path traversal outside `output_dir`
+                if not member_path.resolve().is_relative_to(output_dir.resolve()):
+                    msg = f"Attempted Path Traversal in TAR File: {member.name}"
+                    raise ValueError(msg)
+                bundle.extract(member, path=output_dir)
+
+        # Check and extract based on file type
+        if is_zipfile(bundle_path):
+            with ZipFile(bundle_path, "r") as zip_bundle:
+                _safe_extract_zip(zip_bundle, output_dir)
+        elif tarfile.is_tarfile(bundle_path):
+            with tarfile.open(bundle_path, "r:*") as tar_bundle:
+                _safe_extract_tar(tar_bundle, output_dir)
+        else:
+            msg = f"Unsupported bundle format: {bundle_path.suffix}"
+            raise ValueError(msg)
+
+    def _filter_and_mark_files(self, files: list[BaseFile]) -> list[BaseFile]:
+        """Validate file types and mark files for removal.
+
+        Args:
+            files (list[BaseFile]): List of BaseFile instances.
+
+        Returns:
+            list[BaseFile]: Validated BaseFile instances.
+
+        Raises:
+            ValueError: If unsupported files are encountered and `ignore_unsupported_extensions` is False.
+        """
+        final_files = []
+        ignored_files = []
+
+        for file in files:
+            if not file.path.is_file():
+                self.log(f"Not a file: {file.path.name}")
+                continue
+
+            if file.path.suffix[1:].lower() not in self.valid_extensions:
+                if self.ignore_unsupported_extensions:
+                    ignored_files.append(file.path.name)
+                    continue
+                msg = f"Unsupported file extension: {file.path.suffix}"
+                self.log(msg)
+                if not self.silent_errors:
+                    raise ValueError(msg)
+
+            final_files.append(file)
+
+        if ignored_files:
+            self.log(f"Ignored files: {ignored_files}")
+
+        return final_files