Preface

sorted built-in function has key keyword-only parameter. Most of the custom implementations like the following for heapsort algorithm

import heapq
import itertools
from typing import (Iterable,
                    TypeVar)

Domain = TypeVar('Domain')


def heapsort(iterable: Iterable[Domain]) -> Iterable[Domain]:
    heap = []
    for element in iterable:
        heapq.heappush(heap, element)
    for _ in itertools.repeat(None, len(heap)):
        yield heapq.heappop(heap)

has no key parameter support. I thought of a decorator that will add key parameter support in the next fashion

from functools import wraps
from operator import itemgetter
from typing import (Callable,
                    Iterable,
                    Optional,
                    TypeVar)

Domain = TypeVar('Domain')
Sortable = TypeVar('Sortable')


def with_key(plain: Callable[[Iterable[Sortable]], Iterable[Sortable]]
             ) -> Callable[..., Iterable[Domain]]:
    @wraps(plain)
    def implementation(iterable: Iterable[Domain],
                       *,
                       key: Optional[Callable[[Domain], Sortable]] = None
                       ) -> Iterable[Domain]:
        if key is None:
            yield from plain(iterable)
            return
        yield from map(itemgetter(2),
                       plain((key(element), index, element)
                             for index, element in enumerate(iterable)))

    return implementation

and after that used like

>>> heapsort_with_key = with_key(heapsort)
>>> list(heapsort_with_key(range(10), key=lambda x: -x))
[9, 8, 7, 6, 5, 4, 3, 2, 1, 0]

Problem

I wonder if there is a better way (e.g. in terms of memory efficiency) of adding key parameter? Does anyone know how it is implemented in sorted?