Merge pull request #1032 from mauricesvp/pgn-time-control

Add PGN TimeControl Header Parser, closes #1031

chess/pgn.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import abc
+import dataclasses
 import enum
 import itertools
 import logging
@@ -141,6 +142,39 @@ class SkipType(enum.Enum):
 ResultT = TypeVar("ResultT", covariant=True)
 
 
+class TimeControlType(enum.Enum):
+    UNKNOW = 0
+    UNLIMITED = 1
+    STANDARD = 2
+    RAPID = 3
+    BLITZ = 4
+    BULLET = 5
+
+
+@dataclasses.dataclass
+class TimeControlPart:
+    moves: int = 0
+    time: int = 0
+    increment: float = 0
+    delay: float = 0
+
+
+@dataclasses.dataclass
+class TimeControl:
+    """
+    PGN TimeControl Parser
+    Spec: http://www.saremba.de/chessgml/standards/pgn/pgn-complete.htm#c9.6
+
+    Not Yet Implemented:
+    - Hourglass/Sandclock ('*' prefix)
+    - Differentiating between Bronstein and Simple Delay (Not part of the PGN Spec)
+      - More Info: https://en.wikipedia.org/wiki/Chess_clock#Timing_methods
+    """
+
+    parts: list[TimeControlPart] = dataclasses.field(default_factory=list)
+    type: TimeControlType = TimeControlType.UNKNOW
+
+
 class _AcceptFrame:
     def __init__(self, node: ChildNode, *, is_variation: bool = False, sidelines: bool = True):
         self.state = "pre"
@@ -859,6 +893,14 @@ def accept(self, visitor: BaseVisitor[ResultT]) -> ResultT:
         visitor.end_game()
         return visitor.result()
 
+    def time_control(self) -> TimeControl:
+        """
+        Returns the time control of the game. If the game has no time control
+        information, the default time control ('UNKNOWN') is returned.
+        """
+        time_control_header = self.headers.get("TimeControl", "")
+        return parse_time_control(time_control_header)
+
     @classmethod
     def from_board(cls: Type[GameT], board: chess.Board) -> GameT:
         """Creates a game from the move stack of a :class:`~chess.Board()`."""
@@ -1763,3 +1805,61 @@ def skip_game(handle: TextIO) -> bool:
     Skips a game. Returns ``True`` if a game was found and skipped.
     """
     return bool(read_game(handle, Visitor=SkipVisitor))
+
+
+def parse_time_control(time_control: str) -> TimeControl:
+    tc = TimeControl()
+
+    if not time_control:
+        return tc
+
+    if time_control.startswith("?"):
+        return tc
+
+    if time_control.startswith("-"):
+        tc.type = TimeControlType.UNLIMITED
+        return tc
+
+    def _parse_part(part: str) -> TimeControlPart:
+        tcp = TimeControlPart()
+
+        moves_time, *bonus = part.split("+")
+
+        if bonus:
+            _bonus = bonus[0]
+            if _bonus.lower().endswith("d"):
+                tcp.delay = float(_bonus[:-1])
+            else:
+                tcp.increment = float(_bonus)
+
+        moves, *time = moves_time.split("/")
+        if time:
+            tcp.moves = int(moves)
+            tcp.time = int(time[0])
+        else:
+            tcp.moves = 0
+            tcp.time = int(moves)
+
+        return tcp
+
+    tc.parts = [_parse_part(part) for part in time_control.split(":")]
+
+    if len(tc.parts) > 1:
+        for part in tc.parts[:-1]:
+            if part.moves == 0:
+                raise ValueError("Only last part can be 'sudden death'.")
+
+    # Classification according to https://www.fide.com/FIDE/handbook/LawsOfChess.pdf
+    # (Bullet added)
+    base_time = tc.parts[0].time
+    increment = tc.parts[0].increment
+    if (base_time + 60 * increment) < 3 * 60:
+        tc.type = TimeControlType.BULLET
+    elif (base_time + 60 * increment) < 15 * 60:
+        tc.type = TimeControlType.BLITZ
+    elif (base_time + 60 * increment) < 60 * 60:
+        tc.type = TimeControlType.RAPID
+    else:
+        tc.type = TimeControlType.STANDARD
+
+    return tc

data/pgn/nepomniachtchi-liren-game1.pgn

@@ -0,0 +1,22 @@
+[Event "FIDE World Championship 2023"]
+[Site "Astana KAZ"]
+[Date "2023.04.09"]
+[Round "1"]
+[White "Nepomniachtchi, Ian"]
+[Black "Liren, Ding"]
+[Result "1/2-1/2"]
+[TimeControl "40/7200:20/3600:900+30"]
+[WhiteFideId "4168119"]
+[BlackFideId "8603677"]
+[WhiteElo "2795"]
+[BlackElo "2788"]
+
+1. e4 e5 2. Nf3 Nc6 3. Bb5 a6 4. Ba4 Nf6 5. O-O Be7 6. Bxc6 dxc6
+7. Re1 Nd7 8. d4 exd4 9. Qxd4 O-O 10. Bf4 Nc5 11. Qe3 Bg4 12. Nd4
+Qd7 13. Nc3 Rad8 14. Nf5 Ne6 15. Nxe7+ Qxe7 16. Bg3 Bh5 17. f3 f6
+18. h3 h6 19. Kh2 Bf7 20. Rad1 b6 21. a3 a5 22. Ne2 Rxd1 23. Rxd1
+Rd8 24. Rd3 c5 25. Qd2 c6 26. Rxd8+ Nxd8 27. Qf4 b5 28. Qb8 Kh7
+29. Bd6 Qd7 30. Ng3 Ne6 31. f4 h5 32. c3 c4 33. h4 Qd8 34. Qb7
+Be8 35. Nf5 Qd7 36. Qb8 Qd8 37. Qxd8 Nxd8 38. Nd4 Nb7 39. e5 Kg8
+40. Kg3 Bd7 41. Bc7 Nc5 42. Bxa5 Kf7 43. Bb4 Nd3 44. e6+ Bxe6
+45. Nxc6 Bd7 46. Nd4 Nxb2 47. Kf3 Nd3 48. g3 Nc1 49. Ke3 1/2-1/2

test.py

@@ -2493,6 +2493,53 @@ def test_read_headers(self):
             self.assertEqual(first_drawn_game.headers["Site"], "03")
             self.assertEqual(first_drawn_game[0].move, chess.Move.from_uci("d2d3"))
 
+    def test_parse_time_control(self):
+        with open("data/pgn/nepomniachtchi-liren-game1.pgn") as pgn:
+            game = chess.pgn.read_game(pgn)
+            tc = game.time_control()
+
+            self.assertEqual(tc, chess.pgn.parse_time_control(game.headers["TimeControl"]))
+
+            self.assertEqual(tc.type, chess.pgn.TimeControlType.STANDARD)
+            self.assertEqual(len(tc.parts), 3)
+
+            tcp1, tcp2, tcp3 = tc.parts
+
+            self.assertEqual(tcp1, chess.pgn.TimeControlPart(40, 7200))
+            self.assertEqual(tcp2, chess.pgn.TimeControlPart(20, 3600))
+            self.assertEqual(tcp3, chess.pgn.TimeControlPart(0, 900, 30))
+
+        self.assertEqual(chess.pgn.TimeControlType.BULLET, chess.pgn.parse_time_control("60").type)
+        self.assertEqual(chess.pgn.TimeControlType.BULLET, chess.pgn.parse_time_control("60+1").type)
+
+        self.assertEqual(chess.pgn.TimeControlType.BLITZ, chess.pgn.parse_time_control("60+2").type)
+        self.assertEqual(chess.pgn.TimeControlType.BLITZ, chess.pgn.parse_time_control("300").type)
+        self.assertEqual(chess.pgn.TimeControlType.BLITZ, chess.pgn.parse_time_control("300+3").type)
+
+        self.assertEqual(chess.pgn.TimeControlType.RAPID, chess.pgn.parse_time_control("300+10").type)
+        self.assertEqual(chess.pgn.TimeControlType.RAPID, chess.pgn.parse_time_control("1800").type)
+        self.assertEqual(chess.pgn.TimeControlType.RAPID, chess.pgn.parse_time_control("1800+10").type)
+
+        self.assertEqual(chess.pgn.TimeControlType.STANDARD, chess.pgn.parse_time_control("1800+30").type)
+        self.assertEqual(chess.pgn.TimeControlType.STANDARD, chess.pgn.parse_time_control("5400").type)
+        self.assertEqual(chess.pgn.TimeControlType.STANDARD, chess.pgn.parse_time_control("5400+30").type)
+
+        with self.assertRaises(ValueError):
+            chess.pgn.parse_time_control("300+a")
+
+        with self.assertRaises(ValueError):
+            chess.pgn.parse_time_control("300+ad")
+
+        with self.assertRaises(ValueError):
+            chess.pgn.parse_time_control("600:20/180")
+
+        with self.assertRaises(ValueError):
+            chess.pgn.parse_time_control("abc")
+
+        with self.assertRaises(ValueError):
+            chess.pgn.parse_time_control("40/abc")
+
+
     def test_visit_board(self):
         class TraceVisitor(chess.pgn.BaseVisitor):
             def __init__(self):