File size: 1,922 Bytes
280363b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | from __future__ import annotations
import importlib.util
import shutil
import tempfile
from collections.abc import Callable
from importlib import resources
from pathlib import Path
import chess
WorkspaceEvalFn = Callable[[chess.Board], int]
class WorkspaceManager:
allowed_files = {"eval.py"}
def __init__(self) -> None:
self.root = Path(tempfile.mkdtemp(prefix="zero960_"))
self._write_template_files()
def _write_template_files(self) -> None:
template = resources.files("zero960.workspace_template").joinpath("eval.py").read_text()
self.write_file("eval.py", template)
def cleanup(self) -> None:
shutil.rmtree(self.root, ignore_errors=True)
def resolve_path(self, relative_path: str) -> Path:
if relative_path not in self.allowed_files:
raise ValueError(f"writes are restricted to {sorted(self.allowed_files)}")
return self.root / relative_path
def read_file(self, relative_path: str) -> str:
path = self.resolve_path(relative_path)
return path.read_text()
def write_file(self, relative_path: str, content: str) -> None:
path = self.resolve_path(relative_path)
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(content)
def load_eval_function(self) -> WorkspaceEvalFn:
path = self.resolve_path("eval.py")
module_name = f"zero960_eval_{id(path)}"
spec = importlib.util.spec_from_file_location(module_name, path)
if spec is None or spec.loader is None:
raise RuntimeError("failed to create module spec for eval.py")
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
eval_fn = getattr(module, "evaluate", None)
if eval_fn is None or not callable(eval_fn):
raise RuntimeError("eval.py must define evaluate(board)")
return eval_fn
|