map_except in more_itertools example

map_except_example.py

from typing import Dict, Iterable, Set, Tuple

from more_itertools import map_except


def process(
    names: Iterable[str], whitelisted_names: Set[str], name_to_email: Dict[str, str]
) -> Iterable[str]:
    whitelisted_name_to_email = {
        name: email for name, email in name_to_email.items() if name in whitelisted_names
    }  # 1
    return map_except(
        lambda name: whitelisted_name_to_email[name],  # 3
        names,  # 2
        KeyError  # 4
    )


def test_only_return_emails_for_approved_users():
    name_to_email = {"John": "john.doe@gmail.com", "Bob": "is.still.using@hotmail.com"}
    names = {"Alice", "Bernard", "Bill"} | set(name_to_email.keys())

    actual_emails = process(
        names=names,
        whitelisted_names=set(name_to_email.keys()),
        name_to_email=name_to_email,
    )

    assert set(actual_emails) == set(name_to_email.values())