diff --git a/chess/pgn.py b/chess/pgn.py index 3d4d7e5ff..b51833b8f 100644 --- a/chess/pgn.py +++ b/chess/pgn.py @@ -933,6 +933,40 @@ def time_control(self) -> TimeControl: time_control_header = self.headers.get("TimeControl", "") return parse_time_control(time_control_header) + def set_time_control(self, time_control: TimeControl) -> None: + """ + Sets the ``TimeControl`` header of the game. + + A time control without parts is written as ``-`` if its type is + ``UNLIMITED``, otherwise as ``?`` (unknown). + + :raises: :exc:`ValueError` if the time control cannot be represented + in PGN: only the last part can be sudden death, and a part + cannot have both an increment and a delay. + """ + if not time_control.parts: + self.headers["TimeControl"] = "-" if time_control.type == TimeControlType.UNLIMITED else "?" + return + + def _format_seconds(seconds: float) -> str: + return f"{seconds:.3f}".rstrip("0").rstrip(".") + + def _format_part(part: TimeControlPart) -> str: + text = f"{part.moves}/{part.time}" if part.moves else f"{part.time}" + if part.increment and part.delay: + raise ValueError("Time control part cannot have both increment and delay.") + elif part.increment: + text += f"+{_format_seconds(part.increment)}" + elif part.delay: + text += f"+{_format_seconds(part.delay)}d" + return text + + for part in time_control.parts[:-1]: + if part.moves == 0: + raise ValueError("Only last part can be 'sudden death'.") + + self.headers["TimeControl"] = ":".join(_format_part(part) for part in time_control.parts) + @classmethod def from_board(cls: Type[GameT], board: chess.Board) -> GameT: """Creates a game from the move stack of a :class:`~chess.Board()`.""" diff --git a/docs/pgn.rst b/docs/pgn.rst index dd5a6fa47..108b035a8 100644 --- a/docs/pgn.rst +++ b/docs/pgn.rst @@ -56,7 +56,7 @@ Both extend :class:`~chess.pgn.GameNode`. :members: .. autoclass:: chess.pgn.Game - :members: headers, errors, setup, accept, from_board, without_tag_roster + :members: headers, errors, setup, accept, time_control, set_time_control, from_board, without_tag_roster .. autoclass:: chess.pgn.ChildNode :members: parent, move, starting_comment, nags, san, uci, end diff --git a/test.py b/test.py index a481c2adb..4a5662925 100755 --- a/test.py +++ b/test.py @@ -2627,6 +2627,31 @@ def test_parse_time_control(self): with self.assertRaises(ValueError): chess.pgn.parse_time_control("40/abc") + def test_set_time_control(self): + game = chess.pgn.Game() + + for header in ["40/7200:20/3600:900+30", "300+2", "180+0.5", "5400+30d", "1800", "-", "?"]: + game.set_time_control(chess.pgn.parse_time_control(header)) + self.assertEqual(game.headers["TimeControl"], header) + self.assertEqual(game.time_control(), chess.pgn.parse_time_control(header)) + + game.set_time_control(chess.pgn.TimeControl( + parts=[chess.pgn.TimeControlPart(moves=40, time=6000, increment=30.0), chess.pgn.TimeControlPart(time=900, increment=30)], + type=chess.pgn.TimeControlType.STANDARD)) + self.assertEqual(game.headers["TimeControl"], "40/6000+30:900+30") + + game.set_time_control(chess.pgn.TimeControl()) + self.assertEqual(game.headers["TimeControl"], "?") + self.assertEqual(game.time_control().type, chess.pgn.TimeControlType.UNKNOWN) + + with self.assertRaises(ValueError): + game.set_time_control(chess.pgn.TimeControl(parts=[chess.pgn.TimeControlPart(time=600), chess.pgn.TimeControlPart(moves=20, time=180)])) + + with self.assertRaises(ValueError): + game.set_time_control(chess.pgn.TimeControl(parts=[chess.pgn.TimeControlPart(time=600, increment=5, delay=5)])) + + self.assertEqual(game.headers["TimeControl"], "?") + def test_visit_board(self): class TraceVisitor(chess.pgn.BaseVisitor):