Source code for langchain_community.document_loaders.toml

import json
from pathlib import Path
from typing import Iterator, Union

from langchain_core.documents import Document

from langchain_community.document_loaders.base import BaseLoader


class TomlLoader(BaseLoader):
    """加载`TOML`文件。

    可以加载单个源文件或单个目录中的多个文件。"""

    def __init__(self, source: Union[str, Path]):
        """使用源文件或目录初始化TomlLoader。"""
        self.source = Path(source)

    def lazy_load(self) -> Iterator[Document]:
        """从源文件或目录中延迟加载TOML文档。"""
        import tomli

        if self.source.is_file() and self.source.suffix == ".toml":
            files = [self.source]
        elif self.source.is_dir():
            files = list(self.source.glob("**/*.toml"))
        else:
            raise ValueError("Invalid source path or file type")

        for file_path in files:
            with file_path.open("r", encoding="utf-8") as file:
                content = file.read()
                try:
                    data = tomli.loads(content)
                    doc = Document(
                        page_content=json.dumps(data),
                        metadata={"source": str(file_path)},
                    )
                    yield doc
                except tomli.TOMLDecodeError as e:
                    print(f"Error parsing TOML file {file_path}: {e}")  # noqa: T201