Check Mail

#!/bin/env python
"""
Retrieve number and subjects of unread mail from Gmail.

External dependencies:
    Set a GMAIL_LOGIN environment variable as a json string with Gmail access credentials.
    An example in Zsh:
        export GMAIL_LOGIN='{"username":"[email protected]", "password":"app_password"}'

Usage example:
    In Zsh:
        $ chmod +x checkmail.py
        $ ./checkmail.py

        Or drop the .py and copy the script to $HOME/.local/bin, then:
        $ checkmail
"""

# /// script
# requires-python = ">=3.13"
# dependencies = ["prettytable"]
# ///

import imaplib
import json
import os
from email import message_from_bytes, policy
from socket import gaierror
from typing import Generator, Iterable, cast

try:
    from prettytable import PrettyTable
except ImportError:
    print("Please install the prettytable package.")
    exit(-1)


def check_unread_email(
    *, username: str, password: str, gmail_handle: imaplib.IMAP4_SSL
) -> Generator[dict[str, str], None, None]:
    """Check if there's unread mail in Gmail

    Args:
        username: Gmail email address.
        password: Google app password https://support.google.com/accounts/answer/185833
        gmail_handle: An IMAP SSL object connected to Gmail.

    Returns:
        A generator of string(s).

    Side effect:
        Prints a message if there's no unread mail.

    Raises:
        None
    """

    try:
        gmail_handle.login(username, password)
        gmail_handle.select("INBOX")

        _, response = gmail_handle.search(None, "(UNSEEN)")

        if unread := response[0].split():
            for email_number in reversed(unread):
                status, email_data = gmail_handle.fetch(
                    email_number, "(BODY.PEEK[HEADER.FIELDS (FROM SUBJECT)])"
                )
                if status == "OK":
                    headers = message_from_bytes(
                        cast(bytes, email_data[0][1]),  # pyright: ignore[reportOptionalSubscript]
                        policy=policy.SMTP,
                    )
                    yield {k.capitalize(): v for k, v in headers.items()}
        else:
            print("Inbox zero!")
            exit(0)

    except imaplib.IMAP4.error as e:
        print(e)
        exit(-1)

    finally:
        gmail_handle.close()
        gmail_handle.logout()


def presentation(*, unread_mail: Iterable[dict[str, str]]):
    try:
        table = PrettyTable()
        table.field_names = ("#", "Sender", "Subject")

        for num, headers in enumerate(unread_mail, start=1):
            _from = headers["From"].split("<")[0]
            _from = _from if len(_from) <= 21 else _from[:21] + "..."

            _subject = headers["Subject"]
            _subject = _subject if len(_subject) <= 55 else _subject[:55] + "..."

            table.add_row((num, _from, _subject))

        table.align = "l"
        print(table)

    except (gaierror, ConnectionResetError):
        print("Could not connect to Gmail")


if __name__ == "__main__":
    presentation(
        unread_mail=check_unread_email(
            gmail_handle=imaplib.IMAP4_SSL("imap.gmail.com"),
            **json.loads(os.getenv("GMAIL_LOGIN")),  # pyright:ignore[reportArgumentType]
        )
    )