xref: /aosp_15_r20/external/pytorch/tools/testing/target_determination/determinator.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1from __future__ import annotations
2
3import sys
4from typing import Any
5
6from tools.testing.target_determination.heuristics import (
7    AggregatedHeuristics as AggregatedHeuristics,
8    HEURISTICS,
9    TestPrioritizations as TestPrioritizations,
10)
11
12
13def get_test_prioritizations(
14    tests: list[str], file: Any = sys.stdout
15) -> AggregatedHeuristics:
16    aggregated_results = AggregatedHeuristics(tests)
17    print(f"Received {len(tests)} tests to prioritize", file=file)
18    for test in tests:
19        print(f"  {test}", file=file)
20
21    for heuristic in HEURISTICS:
22        new_rankings: TestPrioritizations = heuristic.get_prediction_confidence(tests)
23        aggregated_results.add_heuristic_results(heuristic, new_rankings)
24
25        print(f"Results from {heuristic.__class__.__name__}")
26        print(new_rankings.get_info_str(verbose=False), file=file)
27
28    return aggregated_results
29