From 98fb86840c9b946ad7a21db390bd4aa93190d182 Mon Sep 17 00:00:00 2001 From: Ben Stull Date: Fri, 5 Jun 2026 03:40:38 -0700 Subject: [PATCH] feat(selection): expose ranked_candidates(); refactor select() onto it Per docs/superpowers/plans/2026-06-04-experience-simulator.md Task 1. Co-Authored-By: Claude Opus 4.8 (1M context) --- hef/selection.py | 52 ++++++++++++++++---- tests/test_selection_ranked.py | 87 ++++++++++++++++++++++++++++++++++ 2 files changed, 129 insertions(+), 10 deletions(-) create mode 100644 tests/test_selection_ranked.py diff --git a/hef/selection.py b/hef/selection.py index afef502..5be1f3f 100644 --- a/hef/selection.py +++ b/hef/selection.py @@ -58,6 +58,38 @@ def candidates_for_mode(records, mode: str, pool_size: int) -> list[Record]: from typing import Optional +def ranked_candidates( + records, + coord: Coordinate, + mode: str, + *, + pool_size: int = 4, + weights: Weights = Weights(), + approved_only: bool = False, +) -> list[tuple[Record, float]]: + """Mode-eligible records sorted nearest-first, each paired with its distance. + + Applies the same approved_only filter and 'av' pool fallback as select(), + then returns up to pool_size nearest (record, distance) pairs. This is the + single source of truth for "the pool"; select() picks from it. + """ + if mode not in CONTENT_MODES: + raise ValueError( + f"ranked_candidates expects a content mode {sorted(CONTENT_MODES)}, " + f"got {mode!r}" + ) + pool = records + if approved_only: + pool = [r for r in pool if r.review_status == "approved"] + candidates = candidates_for_mode(pool, mode, pool_size) + ranked = sorted( + candidates, + key=lambda r: (distance(coord, record_coordinate(r), weights), r.id), + ) + nearest = ranked[:pool_size] + return [(r, distance(coord, record_coordinate(r), weights)) for r in nearest] + + def select( records, coord: Coordinate, @@ -82,17 +114,17 @@ def select( ) if mode == "none": return None - pool = records - if approved_only: - pool = [r for r in pool if r.review_status == "approved"] - candidates = candidates_for_mode(pool, mode, pool_size) - if not candidates: - return None - ranked = sorted( - candidates, - key=lambda r: (distance(coord, record_coordinate(r), weights), r.id), + ranked = ranked_candidates( + records, + coord, + mode, + pool_size=pool_size, + weights=weights, + approved_only=approved_only, ) - nearest = ranked[:pool_size] + if not ranked: + return None + nearest = [r for r, _ in ranked] if rng is None: return nearest[0] return rng.choice(nearest) diff --git a/tests/test_selection_ranked.py b/tests/test_selection_ranked.py new file mode 100644 index 0000000..80da4b9 --- /dev/null +++ b/tests/test_selection_ranked.py @@ -0,0 +1,87 @@ +import random + +import pytest + +from hef.catalog import Record +from hef.selection import Coordinate, Weights, ranked_candidates, select + + +def make_record(**overrides): + base = dict( + id="r", + title="t", + source_url="u", + source_archive="internet_archive", + license="public_domain", + mode="video", + left=0, + right=0, + dark=0, + light=0, + duration_s=600, + file_path="", + ) + base.update(overrides) + return Record(**base) + + +def test_ranked_candidates_sorts_nearest_first_with_distances(): + near = make_record(id="near", mode="video", left=1, right=1, dark=0, light=0) + far = make_record(id="far", mode="video", left=4, right=4, dark=4, light=4) + ranked = ranked_candidates([far, near], Coordinate(0, 0, 0, 0), "video") + assert [r.id for r, _ in ranked] == ["near", "far"] + assert ranked[0][1] < ranked[1][1] + + +def test_ranked_candidates_caps_at_pool_size(): + recs = [make_record(id=f"r{i}", mode="video", left=i % 5) for i in range(10)] + ranked = ranked_candidates(recs, Coordinate(0, 0, 0, 0), "video", pool_size=3) + assert len(ranked) == 3 + + +def test_ranked_candidates_filters_by_mode(): + a = make_record(id="aud", mode="audio") + v = make_record(id="vid", mode="video") + ranked = ranked_candidates([a, v], Coordinate(0, 0, 0, 0), "audio") + assert [r.id for r, _ in ranked] == ["aud"] + + +def test_ranked_candidates_av_falls_back_to_audio_and_video(): + a = make_record(id="aud", mode="audio") + v = make_record(id="vid", mode="video") + ranked = ranked_candidates([a, v], Coordinate(0, 0, 0, 0), "av", pool_size=4) + assert {r.id for r, _ in ranked} == {"aud", "vid"} + + +def test_ranked_candidates_approved_only(): + p = make_record(id="prop", mode="video", review_status="proposed") + ok = make_record(id="appr", mode="video", review_status="approved") + ranked = ranked_candidates([p, ok], Coordinate(0, 0, 0, 0), "video", approved_only=True) + assert [r.id for r, _ in ranked] == ["appr"] + + +def test_ranked_candidates_rejects_none_mode(): + with pytest.raises(ValueError): + ranked_candidates([], Coordinate(0, 0, 0, 0), "none") + + +def test_select_returns_rank_one_of_ranked_candidates(): + recs = [ + make_record(id="near", mode="video", left=1), + make_record(id="far", mode="video", left=4), + ] + coord = Coordinate(0, 0, 0, 0) + ranked = ranked_candidates(recs, coord, "video") + assert select(recs, coord, "video", rng=None).id == ranked[0][0].id + + +def test_select_none_mode_still_returns_none(): + recs = [make_record(id="v", mode="video")] + assert select(recs, Coordinate(0, 0, 0, 0), "none") is None + + +def test_select_shuffles_within_pool_when_rng_given(): + recs = [make_record(id=f"r{i}", mode="video", left=i) for i in range(5)] + coord = Coordinate(0, 0, 0, 0) + picks = {select(recs, coord, "video", pool_size=5, rng=random.Random(s)).id for s in range(20)} + assert len(picks) > 1