inboxos-v1/pipelines/mbox_stream.py

"""Streaming .mbox parser.

A `.mbox` file is a sequence of RFC 5322 messages, each preceded
by a separator line that begins with the literal five characters
`"From "` (capital F, lowercase rom, space). Per RFC 4155 the
canonical separator pattern is a line break followed by `"From "`
so that body lines starting with `">From "` (the mbox escape of
an internal `"From "`) are not mistaken for a new message.

This module yields one parsed (Message, raw_bytes, end_offset)
triple at a time so the historical pipeline can process a multi-GB
mbox without loading the whole file into memory.

Why not `mailbox.mbox`? The stdlib version reads the whole file
into memory and rewrites it on every commit, which is unusable
for the gigabyte-scale Takeout export. This streaming parser
walks the file once and is offset-based, so the pipeline can
checkpoint and resume mid-file.

Offset semantics:
- `end_offset` returned to the caller is the byte position of the
  line break that precedes the *next* separator (i.e. one byte
  past the final byte of the message). On resume we re-scan from
  that exact position; if it lands inside the body of a message
  the next `\nFrom ` pattern picks up the proper boundary.
- For the final message of the file, `end_offset == total`.
"""

from __future__ import annotations

import email
import email.policy
from email.message import Message
from typing import Iterator


# The mbox "From " envelope line. We match a leading newline + "From "
# to avoid false positives on body lines like ">From " (mbox escape).
MBOX_SEP = b"\nFrom "
# A `From ` envelope line that may also be the very first line of
# the file (no preceding newline). Used for the initial scan only.
MBOX_SEP_AT_START = b"From "


def _find_separator(buf: bytes, start: int) -> int:
    """Return the byte index of the next `\nFrom ` at or after `start`.

    The returned index points at the `\n` that precedes the "From "
    envelope line. Callers should skip 5 bytes ("\nFrom ") to land
    on the start of the envelope line itself.
    """
    if start <= 0:
        return buf.find(MBOX_SEP, 0)
    return buf.find(MBOX_SEP, start)


def _first_separator(buf: bytes) -> int:
    """Return the index of the first valid `From ` envelope line.

    A file may begin with a stray "From " line (no preceding newline)
    at offset 0; that is treated as the first separator. If there is
    no leading separator we return -1 and the caller will treat the
    entire file as one message body (last-ditch fallback).
    """
    if buf.startswith(MBOX_SEP_AT_START):
        return 0
    return buf.find(MBOX_SEP, 0)


def stream_messages_from_offset(
    raw: bytes,
    start_offset: int,
) -> Iterator[tuple[Message, bytes, int]]:
    """Yield (message, raw_bytes, end_offset) starting at `start_offset`.

    `end_offset` is the byte offset just past the final byte of the
    message — that is, the position the next resume should start
    scanning from. For the last message in the file this equals
    `len(raw)`.
    """
    total = len(raw)
    if total == 0:
        return

    if start_offset <= 0:
        first_sep = _first_separator(raw)
        if first_sep < 0:
            return
        # Skip the leading "From " envelope line (5 bytes plus any
        # additional content up to the next newline).
        line_end = raw.find(b"\n", first_sep)
        if line_end == -1:
            return
        pos = line_end + 1
    else:
        pos = start_offset
        # If we resumed inside a message body, advance to the next
        # separator so we never parse a truncated message.
        if pos >= total:
            return
        next_sep = _find_separator(raw, pos)
        if next_sep == -1:
            # No further messages.
            return
        # `next_sep` points at the '\n' before the next "From " line.
        # Skip that newline + "From " to land on the envelope line.
        env_start = next_sep + 1
        line_end = raw.find(b"\n", env_start)
        if line_end == -1:
            pos = total
        else:
            pos = line_end + 1

    while pos < total:
        # `pos` is the byte right after a separator's newline, i.e.
        # the first byte of the next message body.
        next_sep = raw.find(MBOX_SEP, pos)
        if next_sep == -1:
            msg_end = total
        else:
            msg_end = next_sep  # includes the leading '\n' of sep

        raw_msg = raw[pos:msg_end]
        try:
            msg = email.message_from_bytes(raw_msg, policy=email.policy.default)
        except Exception:
            # A malformed message must not stall the run. We yield
            # a stub Message so the caller can still advance.
            msg = email.message.Message()

        yield msg, raw_msg, msg_end

        if next_sep == -1:
            return

        # Advance past the separator line for the next iteration.
        env_start = next_sep + 1
        line_end = raw.find(b"\n", env_start)
        if line_end == -1:
            return
        pos = line_end + 1