Decorator to add key parameter support to an existing sort function 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]) -&gt; 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]]
             ) -&gt; Callable[..., Iterable[Domain]]:
    @wraps(plain)
    def implementation(iterable: Iterable[Domain],
                       *,
                       key: Optional[Callable[[Domain], Sortable]] = None
                       ) -&gt; 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

&gt;&gt;&gt; heapsort_with_key = with_key(heapsort)
&gt;&gt;&gt; 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?
Python is open-source, and all the source code in the GitHub repository python/cpython. For a great deal of the built-ins, they can be found in cpython/Python/bltinmodule.c
. In the current revision, the beginning of the sorted function definitions is found at line 2175, and the actual function itself is found at line 2201.

However, I'm not sure that knowing the implementation of sorted can help you too much with optimizing heapsort, because sorted uses Timsort.

I can make a suggestion: decorators will always add overhead, and particularly in a case like this, it's probably better to just make a new version of the function that integrates the key parameter with the original function, though this answer might miss the entire point of your question.