move team stuff to program module

Author: ImogenBits

algobattle/battle.py

@@ -250,8 +250,8 @@ def __get_pydantic_core_schema__(cls, source: Type, handler: GetCoreSchemaHandle
                 case _:
                     return tagged_union_schema(
                         choices={
-                            subclass.model_fields["type"].default: subclass.__pydantic_core_schema__
-                            for subclass in Battle._battle_types.values()
+                            battle.Config.model_fields["type"].default: battle.Config.__pydantic_core_schema__
+                            for battle in Battle._battle_types.values()
                         },
                         discriminator="type",
                     )

algobattle/cli.py

@@ -19,8 +19,8 @@
 from algobattle.battle import Battle, Fight
 from algobattle.match import Match, Ui, AlgobattleConfig
 from algobattle.problem import AnyProblem, Problem
-from algobattle.team import Matchup
 from algobattle.util import Role, RunningTimer, flat_intersperse
+from algobattle.program import Matchup
 
 
 @dataclass

algobattle/match.py

@@ -6,7 +6,8 @@
 from itertools import combinations
 from pathlib import Path
 import tomllib
-from typing import Annotated, Any, ClassVar, Self, TypeAlias, TypeVar, TypedDict, cast, overload
+from typing import Annotated, Any, ClassVar, Self, TypeAlias, TypeVar, cast, overload
+from typing_extensions import TypedDict
 
 from pydantic import AfterValidator, ByteSize, ConfigDict, Field, GetCoreSchemaHandler, ValidationInfo, model_validator
 from pydantic.types import PathType
@@ -17,8 +18,7 @@
 from docker.types import LogConfig, Ulimit
 
 from algobattle.battle import Battle, FightHandler, FightUi, BattleUi, Iterated
-from algobattle.program import ProgramUi
-from algobattle.team import BuildUi, Matchup, Team, TeamHandler
+from algobattle.program import ProgramConfigView, ProgramUi, BuildUi, Matchup, Team, TeamHandler
 from algobattle.problem import InstanceT, Problem, ProblemName, SolutionT
 from algobattle.util import (
     ExceptionInfo,
@@ -98,9 +98,7 @@ async def run(
         if ui is None:
             ui = Ui()
 
-        with await TeamHandler.build(
-            config.teams, problem, config.execution.mode, config.match, config.docker, ui
-        ) as teams:
+        with await TeamHandler.build(config.teams, problem, config.as_prog_config(), ui) as teams:
             result = cls(
                 active_teams=[t.name for t in teams.active],
                 excluded_teams=teams.excluded,
@@ -631,3 +629,16 @@ def from_file(cls, file: Path) -> Self:
                 except tomllib.TOMLDecodeError as e:
                     raise ValueError(f"The config file at {file} is not a properly formatted TOML file!\n{e}")
         return cls.model_validate(config_dict, context={"base_path": file.parent})
+
+    def as_prog_config(self) -> ProgramConfigView:
+        """Builds a simple object containing all program relevant settings."""
+        return ProgramConfigView(
+            build_timeout=self.match.build_timeout,
+            max_image_size=self.match.image_size,
+            strict_timeouts=self.match.strict_timeouts,
+            build_kwargs=self.docker.build.kwargs,
+            run_kwargs=self.docker.run.kwargs,
+            generator=self.match.generator,
+            solver=self.match.solver,
+            mode=self.execution.mode,
+        )

algobattle/program.py

@@ -1,17 +1,18 @@
 """Module providing an interface to interact with the teams' programs."""
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
+from itertools import combinations
 from os import environ
 from pathlib import Path
 from tarfile import TarFile, is_tarfile
 from tempfile import TemporaryDirectory
 from timeit import default_timer
 from types import EllipsisType
-from typing import Any, ClassVar, Iterator, Protocol, Self, TypeVar, cast, Generator as PyGenerator
+from typing import Any, ClassVar, Iterator, Mapping, Protocol, Self, TypeVar, cast, Generator as PyGenerator
 from typing_extensions import TypedDict
 from uuid import uuid4
 import json
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from zipfile import ZipFile, is_zipfile
 
 from docker import DockerClient
@@ -32,6 +33,7 @@
     ExceptionInfo,
     ExecutionError,
     ExecutionTimeout,
+    MatchMode,
     ValidationError,
     Role,
     BaseModel,
@@ -79,36 +81,13 @@ class RunSpecs:
     overriden: RunConfigOverride
 
 
-@dataclass(frozen=True, slots=True)
-class RunConfigView:
+class RunConfigView(Protocol):
     """Config view for single runs."""
 
     timeout: float | None
     space: int | None
     cpus: int
 
-    def reify(
-        self,
-        timeout: float | None | EllipsisType,
-        space: int | None | EllipsisType,
-        cpus: int | EllipsisType,
-    ) -> RunSpecs:
-        """Merges the overriden config options with the parsed ones."""
-        overriden = RunConfigOverride()
-        if timeout is ...:
-            timeout = self.timeout
-        else:
-            overriden["timeout"] = timeout
-        if space is ...:
-            space = self.space
-        else:
-            overriden["space"] = space
-        if cpus is ...:
-            cpus = self.cpus
-        else:
-            overriden["cpus"] = cpus
-        return RunSpecs(timeout=timeout, space=space, cpus=cpus, overriden=overriden)
-
 
 @dataclass(frozen=True, slots=True)
 class ProgramConfigView:
@@ -121,6 +100,7 @@ class ProgramConfigView:
     run_kwargs: dict[str, Any]
     generator: RunConfigView
     solver: RunConfigView
+    mode: MatchMode
 
 
 class ProgramUi(Protocol):
@@ -691,3 +671,177 @@ async def run(
                 battle_data=battle_data,
                 solution=solution,
             )
+
+
+class BuildUi(Protocol):
+    """Provides and interface for the build process to update the ui."""
+
+    @abstractmethod
+    def start_build(self, team: str, role: Role, timeout: float | None) -> None:
+        """Informs the ui that a new program is being built."""
+
+    @abstractmethod
+    def finish_build(self) -> None:
+        """Informs the ui that the current build has been finished."""
+
+
+class _TeamInfo(Protocol):
+    generator: Path
+    solver: Path
+
+
+@dataclass(frozen=True, slots=True)
+class Team:
+    """Class bundling together the programs of a team."""
+
+    name: str
+    generator: Generator
+    solver: Solver
+
+    @classmethod
+    async def build(
+        cls,
+        name: str,
+        info: _TeamInfo,
+        problem: AnyProblem,
+        config: ProgramConfigView,
+        ui: BuildUi,
+    ) -> "Team":
+        """Builds the specified docker files into images and return the corresponding team.
+
+        Args:
+            name: Name of the team.
+            info: Team info containing the paths to the program data.
+            problem: The problem class the current match is fought over.
+            config: Config for the programs.
+
+        Returns:
+            The built team.
+
+        Raises:
+            ValueError: If the team name is already in use.
+            DockerError: If the docker build fails for some reason
+        """
+        tag_name = name.lower().replace(" ", "_") if config.mode == "testing" else None
+        ui.start_build(name, Role.generator, config.build_timeout)
+        generator = await Generator.build(
+            path=info.generator,
+            problem=problem,
+            config=config,
+            team_name=tag_name,
+        )
+        ui.finish_build()
+        try:
+            ui.start_build(name, Role.solver, config.build_timeout)
+            solver = await Solver.build(
+                path=info.solver,
+                problem=problem,
+                config=config,
+                team_name=tag_name,
+            )
+            ui.finish_build()
+        except Exception:
+            generator.remove()
+            raise
+        return Team(name, generator, solver)
+
+    def __str__(self) -> str:
+        return self.name
+
+    def __eq__(self, o: object) -> bool:
+        if isinstance(o, Team):
+            return self.name == o.name
+        else:
+            return False
+
+    def __hash__(self) -> int:
+        return hash(self.name)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, _type: Any, _value: Any, _traceback: Any):
+        self.cleanup()
+
+    def cleanup(self) -> None:
+        """Removes the built docker images."""
+        self.generator.remove()
+        self.solver.remove()
+
+
+@dataclass(frozen=True)
+class Matchup:
+    """Represents an individual matchup of teams."""
+
+    generator: Team
+    solver: Team
+
+    def __iter__(self) -> Iterator[Team]:
+        yield self.generator
+        yield self.solver
+
+    def __repr__(self) -> str:
+        return f"Matchup({self.generator.name}, {self.solver.name})"
+
+
+@dataclass
+class TeamHandler:
+    """Handles building teams and cleaning them up."""
+
+    active: list[Team] = field(default_factory=list)
+    excluded: dict[str, ExceptionInfo] = field(default_factory=dict)
+    cleanup: bool = True
+
+    @classmethod
+    async def build(
+        cls,
+        infos: Mapping[str, _TeamInfo],
+        problem: AnyProblem,
+        config: ProgramConfigView,
+        ui: BuildUi,
+    ) -> Self:
+        """Builds the programs of every team.
+
+        Attempts to build the programs of every team. If any build fails, that team will be excluded and all its
+        programs cleaned up.
+
+        Args:
+            infos: Teams that participate in the match.
+            problem: Problem class that the match will be fought with.
+            config: Config options.
+
+        Returns:
+            :class:`TeamHandler` containing the info about the participating teams.
+        """
+        handler = cls(cleanup=config.mode == "tournament")
+        for name, info in infos.items():
+            try:
+                team = await Team.build(name, info, problem, config, ui)
+                handler.active.append(team)
+            except Exception as e:
+                handler.excluded[name] = ExceptionInfo.from_exception(e)
+        return handler
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(self, _type: Any, _value: Any, _traceback: Any):
+        if self.cleanup:
+            for team in self.active:
+                team.cleanup()
+
+    @property
+    def grouped_matchups(self) -> list[tuple[Matchup, Matchup]]:
+        """All matchups, grouped by the involved teams.
+
+        Each tuple's first matchup has the first team in the group generating, the second has it solving.
+        """
+        return [(Matchup(*g), Matchup(*g[::-1])) for g in combinations(self.active, 2)]
+
+    @property
+    def matchups(self) -> list[Matchup]:
+        """All matchups that will be fought."""
+        if len(self.active) == 1:
+            return [Matchup(self.active[0], self.active[0])]
+        else:
+            return [m for pair in self.grouped_matchups for m in pair]