Add chess.pgn.Game.set_time_control()

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()`."""

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

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):