shows how to use the more_itertools partition function

partition_example.py

from typing import Dict, Iterable, Set, Tuple

from more_itertools import partition


def process(
    names: Iterable[str], whitelisted_names: Set[str], name_to_email: Dict[str, str]
) -> Tuple[Iterable[str], Iterable[str]]:
    refused_names, approved_names = partition(
        lambda name: name in whitelisted_names, names
    )  # 1
    approved_emails = {name_to_email[name] for name in approved_names}  # 2
    return refused_names, approved_emails  # 3


def test_only_return_emails_for_approved_users():
    name_to_email = {"John": "[email protected]", "Bob": "[email protected]"}
    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())  # 4


def test_filters_non_approved_users():
    not_in_whitelist = {"Alice", "Bernard", "Bill"}
    name_to_email = {"John": "[email protected]", "Bob": "[email protected]"}
    names = not_in_whitelist | set(name_to_email.keys())

    refused_names, _ = process(
        names=names,
        whitelisted_names=set(name_to_email.keys()),
        name_to_email=name_to_email,
    )

    assert set(refused_names) == not_in_whitelist  # 5

[email protected]