WalisonCruz/function-gemma

1from dataclasses import dataclass, field2 3try:4    import trackio.utils as utils5except ImportError:6    import utils7 8 9@dataclass10class RunSelection:11    choices: list[str] = field(default_factory=list)12    selected: list[str] = field(default_factory=list)13    locked: bool = False14 15    def update_choices(16        self, runs: list[str], preferred: list[str] | None = None17    ) -> bool:18        if self.choices == runs:19            return False20        new_choices = set(runs) - set(self.choices)21        self.choices = list(runs)22        if self.locked:23            base = set(self.selected) | new_choices24        elif preferred:25            base = set(preferred)26        else:27            base = set(runs)28        self.selected = [run for run in self.choices if run in base]29        return True30 31    def select(self, runs: list[str]) -> list[str]:32        choice_set = set(self.choices)33        self.selected = [run for run in runs if run in choice_set]34        self.locked = True35        return self.selected36 37    def replace_group(38        self, group_runs: list[str], new_subset: list[str] | None39    ) -> tuple[list[str], list[str]]:40        new_subset = utils.ordered_subset(group_runs, new_subset)41        selection_set = set(self.selected)42        selection_set.difference_update(group_runs)43        selection_set.update(new_subset)44        self.selected = [run for run in self.choices if run in selection_set]45        self.locked = True46        return new_subset, self.selected47