diff --git a/hef/selection.py b/hef/selection.py index 3a8d5f3..afef502 100644 --- a/hef/selection.py +++ b/hef/selection.py @@ -53,3 +53,46 @@ def candidates_for_mode(records, mode: str, pool_size: int) -> list[Record]: extra = [r for r in records if r.mode in {"audio", "video"}] return primary + extra return primary + + +from typing import Optional + + +def select( + records, + coord: Coordinate, + mode: str, + *, + pool_size: int = 4, + weights: Weights = Weights(), + approved_only: bool = False, + rng=None, +) -> Optional[Record]: + """Pick the record nearest `coord` for the given selector `mode`. + + - 'none' selector mode returns None (the void/rest state). + - With rng=None, returns the single nearest record (ties broken by id) for + deterministic behavior. Pass a random.Random to shuffle within the + pool_size nearest records. + - approved_only restricts to records the human has blessed. + """ + if mode not in SELECTOR_MODES: + raise ValueError( + f"invalid selector mode {mode!r}; expected one of {sorted(SELECTOR_MODES)}" + ) + 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), + ) + nearest = ranked[:pool_size] + if rng is None: + return nearest[0] + return rng.choice(nearest) diff --git a/tests/test_selection.py b/tests/test_selection.py index 8012986..e535c7e 100644 --- a/tests/test_selection.py +++ b/tests/test_selection.py @@ -79,3 +79,58 @@ def test_av_falls_back_to_audio_and_video_when_thin(): def test_candidates_rejects_none_mode(): with pytest.raises(ValueError): candidates_for_mode([], "none", pool_size=4) + + +import random + +from hef.selection import select + + +def test_none_mode_returns_none(): + records = [make_record(id="v", mode="video")] + assert select(records, Coordinate(0, 0, 0, 0), "none") is None + + +def test_empty_pool_returns_none(): + assert select([], Coordinate(0, 0, 0, 0), "video") is None + + +def test_select_returns_nearest_by_default(): + near = make_record(id="near", mode="video", left=2, right=2, dark=2, light=2) + far = make_record(id="far", mode="video", left=0, right=0, dark=0, light=0) + result = select([far, near], Coordinate(2, 2, 2, 2), "video") + assert result.id == "near" + + +def test_select_is_deterministic_without_rng(): + records = [ + make_record(id="b", mode="video", left=1), + make_record(id="a", mode="video", left=1), + ] + # Equal distance -> tie broken by id, so "a" wins deterministically. + result = select(records, Coordinate(1, 0, 0, 0), "video") + assert result.id == "a" + + +def test_select_with_rng_picks_within_nearest_pool(): + records = [make_record(id=f"r{i}", mode="video", left=i % 5) for i in range(10)] + coord = Coordinate(2, 0, 0, 0) + rng = random.Random(0) + chosen = {select(records, coord, "video", pool_size=3, rng=rng).id for _ in range(50)} + # Only records inside the 3-nearest pool may ever be chosen. + ranked = sorted(records, key=lambda r: abs(r.left - 2)) + allowed = {r.id for r in ranked[:3]} + assert chosen <= allowed + + +def test_approved_only_excludes_proposed(): + proposed = make_record(id="p", mode="video", review_status="proposed") + approved = make_record(id="ok", mode="video", review_status="approved") + result = select([proposed, approved], Coordinate(0, 0, 0, 0), "video", + approved_only=True) + assert result.id == "ok" + + +def test_invalid_selector_mode_raises(): + with pytest.raises(ValueError): + select([], Coordinate(0, 0, 0, 0), "telepathy")