Group mailbox messages through reference identifiers

#!/usr/bin/env python3
'Group mailbox messages through reference identifiers'
import argparse
import json
from pathlib import Path

def write_json(path, value):
    data = json.dumps(value, indent=2, ensure_ascii=False, allow_nan=False) + "\n"
    with Path(path).open("x", encoding="utf-8", newline="") as handle:
        handle.write(data)

import re
import hashlib
import mimetypes
import mailbox
from collections import Counter, defaultdict
from datetime import datetime, timezone
from email import policy
from email.parser import BytesParser
from email.message import EmailMessage
from email.utils import getaddresses, parseaddr, formataddr, parsedate_to_datetime, format_datetime, make_msgid

def raw(path,limit=25*1024*1024):
    if Path(path).stat().st_size>limit: raise ValueError("Input exceeds the size limit.")
    return Path(path).read_bytes()
def message(path):
    return BytesParser(policy=policy.default).parsebytes(raw(path))
def messages(path):
    if Path(path).stat().st_size>100*1024*1024: raise ValueError("Mailbox exceeds 100 MiB.")
    box=mailbox.mbox(path,factory=lambda f:BytesParser(policy=policy.default).parse(f),create=False)
    try:
        result=[]
        for msg in box:
            if len(result)>=5000: raise ValueError("Mailbox exceeds 5000 messages.")
            result.append(msg)
        return result
    finally: box.close()
def text_body(msg):
    part=msg.get_body(preferencelist=("plain",))
    return part.get_content() if part else ""
def address_pairs(msg,headers=("from","to","cc","bcc","reply-to")):
    return getaddresses([str(v) for key in headers for v in msg.get_all(key,[])])
def attachments(msg):
    return [p for p in msg.walk() if p.get_content_maintype()!="multipart" and (p.get_content_disposition()=="attachment" or p.get_filename() is not None)]
def payload(part):
    if part.get_content_type()=="message/rfc822" and isinstance(part.get_payload(),list):
        return part.get_payload(0).as_bytes(policy=policy.default)
    return part.get_payload(decode=True) or b""
def save_message(path,msg):
    data=msg.as_bytes(policy=policy.SMTP)
    with Path(path).open("xb") as handle: handle.write(data)
def clean_address(value):
    if "\r" in value or "\n" in value: raise ValueError("Address cannot contain newlines.")
    display,address=parseaddr(value)
    if "@" not in address or any(c.isspace() for c in address): raise ValueError("Supply one ordinary email address.")
    return address
def draft(sender,recipient,subject,body):
    msg=EmailMessage()
    msg["From"]=clean_address(sender)
    msg["To"]=clean_address(recipient)
    msg["Subject"]=subject
    msg["Date"]=format_datetime(datetime.now(timezone.utc))
    msg["Message-ID"]=make_msgid(domain="scriptage.local")
    msg.set_content(body)
    return msg

def main():
    parser = argparse.ArgumentParser(description='Group mailbox messages through reference identifiers')
    parser.add_argument("--input",required=True,help="Local EML, MBOX, or body file as documented")
    parser.add_argument("--output",required=True,help="New result file or directory")
    args = parser.parse_args()
    parent={}
    def root(key):
        parent.setdefault(key,key)
        while parent[key]!=key:
            parent[key]=parent[parent[key]]
            key=parent[key]
        return key
    items=[]
    for i,msg in enumerate(messages(args.input),1):
        ids=re.findall(r"<[^<>\s]+>",str(msg.get("Message-ID","")))
        own=ids[0] if ids else f"<local-{i}>"
        refs=re.findall(r"<[^<>\s]+>"," ".join(str(v) for h in ("References","In-Reply-To") for v in msg.get_all(h,[])))
        root(own)
        for ref in refs: parent[root(own)]=root(ref)
        items.append((i,own,str(msg.get("Subject",""))))
    groups=defaultdict(list)
    for i,key,subject in items: groups[root(key)].append({"message":i,"subject":subject})
    write_json(args.output,[{"reference_group":key,"messages":group} for key,group in groups.items()])

if __name__ == "__main__":
    try:
        main()
    except (OSError, ValueError, KeyError, TypeError, ArithmeticError, IndexError, LookupError) as exc:
        raise SystemExit(str(exc))