document_loader.py

"""
文档加载器：扫描指定目录，读取文档文件，返回统一格式的 Document 对象
"""

# ============ 导入模块 ============
# Python 用 import 导入模块，相当于 C 的 #include
import os                          # os 模块：操作系统相关，遍历目录、拼接路径等
from dataclasses import dataclass  # 从 dataclasses 模块导入 dataclass 装饰器


# ============ 定义数据类 ============
# @dataclass 是装饰器，自动帮你生成 __init__（构造函数）
# 相当于 C 的 struct Document { char* content; map metadata; };
@dataclass
class Document:
    content: str       # 文件的文本内容（str = 字符串，相当于 std::string）
    metadata: dict     # 元数据（dict = 字典，相当于 std::map<string, string>）


# ============ 支持的文件类型 ============
# Python 里 set 用花括号，类似 C++ 的 std::set
# 只处理这些后缀的文件，其他的跳过
SUPPORTED_EXTENSIONS = {".md", ".txt"}


# ============ 核心函数 ============
# def = 定义函数（相当于 C 的函数声明）
# directory: str = 参数类型提示（Python 不强制，但写了更清晰）
# -> list[Document] = 返回值类型提示（返回一个 Document 的列表）
def load_documents(directory: str) -> list[Document]:
    """
    扫描指定目录，加载所有支持的文档文件。

    参数：
        directory: 要扫描的目录路径，如 "F:/my_documents"

    返回：
        Document 对象的列表
    """
    # [] 是空列表，相当于 C++ 的 vector<Document>
    documents = []

    # 这些目录名要跳过（虚拟环境、隐藏目录、缓存等）
    # Python 里 set 可以直接用 in 判断是否包含，O(1) 查找
    skip_dirs = {".venv", "__pycache__", ".git", "node_modules"}

    # os.walk() 遍历目录树（包括子目录）
    # 它每次返回三个值：当前目录路径、子目录列表、文件列表
    # 相当于 C 的 opendir() + readdir()，但自动递归子目录
    for current_dir, sub_dirs, files in os.walk(directory):

        # 原地修改 sub_dirs 可以阻止 os.walk 进入这些子目录
        # 这是 os.walk 的一个特殊用法：修改 sub_dirs 影响后续遍历
        sub_dirs[:] = [d for d in sub_dirs if d not in skip_dirs]

        # 遍历当前目录下的所有文件名
        for filename in files:

            # os.path.splitext("readme.md") -> ("readme", ".md")
            # [1] 取第二个元素（后缀），Python 下标从 0 开始，跟 C 一样
            # .lower() 转小写，防止 .MD 和 .md 不匹配
            ext = os.path.splitext(filename)[1].lower()

            # not in = 不在集合里，跳过不支持的文件类型
            if ext not in SUPPORTED_EXTENSIONS:
                continue  # 相当于 C 的 continue

            # os.path.join() 拼接路径，自动处理 / 和 \
            # 比手动拼 current_dir + "/" + filename 更安全
            filepath = os.path.join(current_dir, filename)

            # try-except 相当于 C++ 的 try-catch
            # 读文件可能出错（编码问题、权限问题），用 try 保护
            try:
                # with open() as f: 打开文件，出了 with 块自动关闭
                # "r" = 只读模式，encoding="utf-8" 指定编码
                with open(filepath, "r", encoding="utf-8") as f:
                    content = f.read()  # 一次读完整个文件内容
            except Exception as e:
                # f"..." 是格式化字符串，{} 里的变量会被替换成值
                # 相当于 C 的 printf("跳过 %s: %s", filepath, e)
                print(f"跳过 {filepath}: {e}")
                continue

            # 跳过空文件（strip() 去掉首尾空白后判断）
            if not content.strip():
                continue

            # 创建 Document 对象并添加到列表
            # 相当于 C++: documents.push_back(Document{content, metadata})
            doc = Document(
                content=content,
                metadata={
                    "filename": filename,                    # 文件名
                    "filepath": filepath,                    # 完整路径
                    "file_type": ext,                        # 文件类型后缀
                }
            )
            documents.append(doc)  # append = push_back

    print(f"加载完成：共 {len(documents)} 个文档")  # len() = .size()
    return documents


# ============ 测试代码 ============
# __name__ == "__main__" 表示"这个文件被直接运行时"才执行下面的代码
# 如果是被别的文件 import 的，就不执行
# 相当于 C 的 main() 函数
if __name__ == "__main__":
    # 测试：加载当前项目目录下的文档
    docs = load_documents("F:/AI_Program/enterprise-kb")
    for doc in docs:
        # [:100] 是切片，取前 100 个字符，防止打印太长
        print(f"文件: {doc.metadata['filename']}")
        print(f"内容预览: {doc.content[:100]}")
        print("---")