impl/python.py

from functools import cmp_to_key
from typing import Any, List, Mapping, Sequence

from .collections_sort_by_types import SortDirection


# The type ranks that give mixed values a total order. Absent sorts last in
# both directions, so it is ranked above every present value and then excluded
# from the direction flip below.
RANK_BOOL = 0
RANK_NUMBER = 1
RANK_STRING = 2
RANK_ABSENT = 3


def _rank_of(value: Any, key: str) -> int:
    # Absent and null are the same thing: a document that omits a field and one
    # that nulls it mean the same to every reader.
    if value is None:
        return RANK_ABSENT
    # bool before int: in Python True is an int, and ranking it as a number
    # would order it against 0 and 1 instead of with the other booleans.
    if isinstance(value, bool):
        return RANK_BOOL
    if isinstance(value, (int, float)):
        return RANK_NUMBER
    if isinstance(value, str):
        return RANK_STRING
    raise TypeError('cannot sort by the list or map at "%s"' % (key,))


def _compare_scalars(rank: int, a: Any, b: Any) -> int:
    if rank == RANK_STRING:
        # Python compares strings by code point, which is the order this
        # capability pins; TypeScript spells the same comparison out by hand
        # because JavaScript would otherwise compare UTF-16 code units.
        return -1 if a < b else (1 if a > b else 0)
    if rank == RANK_NUMBER:
        return -1 if a < b else (1 if a > b else 0)
    return 0 if a == b else (1 if a else -1)


def sort_by(
    records: Sequence[Mapping[str, Any]],
    key: str,
    direction: SortDirection = "asc",
) -> List[Mapping[str, Any]]:
    """Stably sort ``records`` by ``key``, ascending or descending.

    Ties keep their input order in both directions, which is what lets a user
    sort by one column and then another and get the multi-column sort they
    expect rather than a reshuffle.
    """
    if isinstance(records, (str, bytes)) or not isinstance(records, (list, tuple)):
        raise TypeError("sort_by needs a list of records")
    if not isinstance(key, str) or key == "":
        raise TypeError("sort_by needs a non-empty key name")
    if direction not in ("asc", "desc"):
        raise ValueError('direction must be "asc" or "desc", received "%s"' % (direction,))

    # Rank every record up front. Raising mid-comparison would make the error
    # depend on which comparisons this interpreter's sort happened to perform.
    values = [record.get(key) if isinstance(record, dict) else None for record in records]
    ranks = [_rank_of(value, key) for value in values]
    sign = -1 if direction == "desc" else 1

    indexed = list(range(len(records)))

    def compare(i: int, j: int) -> int:
        ri, rj = ranks[i], ranks[j]
        # Absent values sink to the bottom whichever way the sort runs: nobody
        # wants the rows they know nothing about at the top of a descending table.
        if ri == RANK_ABSENT or rj == RANK_ABSENT:
            if ri == rj:
                return 0
            return 1 if ri == RANK_ABSENT else -1
        if ri != rj:
            base = -1 if ri < rj else 1
        else:
            base = _compare_scalars(ri, values[i], values[j])
        return base * sign

    indexed.sort(key=cmp_to_key(compare))
    return [records[i] for i in indexed]