llama-index

1
"""Simple reader for mbox (mailbox) files."""
2

3
import os
4
from pathlib import Path
5
from typing import Any, List
6

7
from llama_index.legacy.readers.base import BaseReader
8
from llama_index.legacy.readers.file.mbox_reader import MboxReader as MboxFileReader
9
from llama_index.legacy.schema import Document
10

11

12
class MboxReader(BaseReader):
13
    """Mbox e-mail reader.
14

15
    Reads a set of e-mails saved in the mbox format.
16
    """
17

18
    def __init__(self) -> None:
19
        """Initialize."""
20

21
    def load_data(self, input_dir: str, **load_kwargs: Any) -> List[Document]:
22
        """Load data from the input directory.
23

24
        load_kwargs:
25
            max_count (int): Maximum amount of messages to read.
26
            message_format (str): Message format overriding default.
27
        """
28
        docs: List[Document] = []
29
        for dirpath, dirnames, filenames in os.walk(input_dir):
30
            dirnames[:] = [d for d in dirnames if not d.startswith(".")]
31
            for filename in filenames:
32
                if filename.endswith(".mbox"):
33
                    filepath = os.path.join(dirpath, filename)
34
                    file_docs = MboxFileReader(**load_kwargs).load_data(Path(filepath))
35
                    docs.extend(file_docs)
36
        return docs
37