from typing import MutableMapping
from typing import Optional
from typing import Container


def sort_locs_by_size(loc_to_size: MutableMapping[tuple[str, str], int],
                      prioritize: Optional[Container[str]] = None
                      ) -> list[tuple[str, str]]:
    """Sorts locations by size, optionally prioritizing locations from certain
    datasets first.

    Args:
        loc_to_size: dict, maps each (dataset, location) tuple to its size,
            modified in-place
        prioritize: optional list of str, datasets to prioritize

    Returns: list of (dataset, location) tuples, ordered from smallest size to
        largest. Locations from prioritized datasets come first.
    """
    result = []
    if prioritize is not None:
        # modify loc_to_size in place, so copy its keys before iterating
        prioritized_loc_to_size = {
            loc: loc_to_size.pop(loc) for loc in list(loc_to_size.keys())
            if loc[0] in prioritize
        }
        result = sort_locs_by_size(prioritized_loc_to_size)

    result += sorted(loc_to_size, key=loc_to_size.__getitem__)
    return result