Create class to hold multiple comments

chess/pgn.py

@@ -184,6 +184,70 @@ def __init__(self, node: ChildNode, *, is_variation: bool = False, sidelines: bo
         self.in_variation = False
 
 
+class GameNodeComment:
+    """
+    PGN Comment Storage
+
+    A class that can hold one or more comments for a GameNode.
+    """
+    def __init__(self, comment: Union[str, list[str]] = ""):
+        self.set(comment)
+
+    def set(self, new_comment: Union[str, list[str]]):
+        """Replace the comment with a new comment or a list of comments."""
+        self._comments = new_comment if isinstance(new_comment, list) else [new_comment] if new_comment else []
+
+    def pgn_format(self) -> str:
+        """Create a string representation of the comments in PGN format."""
+        comments = list(map(lambda s: s.replace("{", ""), self._comments))
+        comments = list(map(lambda s: s.replace("}", "").strip(), comments))
+        output_comments = GameNodeComment(comments)
+        output_comments.remove_empty()
+        return "{ " + output_comments.join(" } { ") + " }"
+
+    def remove_empty(self) -> None:
+        """Remove empty comments from the comment list."""
+        self._comments = list(filter(None, self._comments))
+
+    def append(self, new_comment: str) -> None:
+        """Append a new comment to the end of the comment list."""
+        self._comments.append(new_comment)
+
+    def extend(self, new_comments: list[str]) -> None:
+        """Append several new comments to the end of the comment list."""
+        self._comments.extend(new_comments)
+
+    def insert(self, index: int, new_comment: str) -> None:
+        """Insert a new comment before the specified index."""
+        self._comments.insert(index, new_comment)
+
+    def join(self, joiner: str) -> str:
+        """Join all of the comments together with a joiner string between each."""
+        return joiner.join(self._comments)
+
+    def __len__(self) -> int:
+        return len(self._comments)
+
+    def __getitem__(self, index: int) -> str:
+        return self._comments[index]
+
+    def __setitem__(self, index: int, new_comment: str) -> None:
+        self._comments[index] = new_comment
+
+    def __add__(self, other: GameNodeComment) -> GameNodeComment:
+        return GameNodeComment(self._comments + other._comments)
+
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, str):
+            return len(self) == 1 and self[0] == other
+        elif isinstance(other, list):
+            return self._comments == other
+        elif isinstance(other, GameNodeComment):
+            return self._comments == other._comments
+        else:
+            return False
+
+
 class GameNode(abc.ABC):
     parent: Optional[GameNode]
     """The parent node or ``None`` if this is the root node of the game."""
@@ -197,24 +261,24 @@ class GameNode(abc.ABC):
     variations: List[ChildNode]
     """A list of child nodes."""
 
-    comment: list[str]
+    comment: GameNodeComment
     """
     A comment that goes behind the move leading to this node. Comments
     that occur before any moves are assigned to the root node.
     """
 
-    starting_comment: list[str]
+    starting_comment: GameNodeComment
     nags: Set[int]
 
     def __init__(self, *, comment: Union[str, list[str]] = "") -> None:
         self.parent = None
         self.move = None
         self.variations = []
-        self.comment = comment if isinstance(comment, list) else [comment] if comment else []
+        self.comment = GameNodeComment(comment)
 
         # Deprecated: These should be properties of ChildNode, but need to
         # remain here for backwards compatibility.
-        self.starting_comment: list[str] = []
+        self.starting_comment = GameNodeComment()
         self.nags = set()
 
     @abc.abstractmethod
@@ -436,7 +500,7 @@ def add_line(self, moves: Iterable[chess.Move], *, comment: Union[str, list[str]
             else:
                 node.comment.extend(comment)
         else:
-            node.comment = comment if isinstance(comment, list) else [comment] if comment else []
+            node.comment.set(comment)
 
         node.nags.update(nags)
 
@@ -449,7 +513,7 @@ def eval(self) -> Optional[chess.engine.PovScore]:
 
         Complexity is `O(n)`.
         """
-        match = EVAL_REGEX.search(" ".join(self.comment))
+        match = EVAL_REGEX.search(self.comment.join(" "))
         if not match:
             return None
 
@@ -475,7 +539,7 @@ def eval_depth(self) -> Optional[int]:
 
         Complexity is `O(1)`.
         """
-        match = EVAL_REGEX.search(" ".join(self.comment))
+        match = EVAL_REGEX.search(self.comment.join(" "))
         return int(match.group("depth")) if match and match.group("depth") else None
 
     def set_eval(self, score: Optional[chess.engine.PovScore], depth: Optional[int] = None) -> None:
@@ -498,7 +562,7 @@ def set_eval(self, score: Optional[chess.engine.PovScore], depth: Optional[int]
             if found:
                 break
 
-        self.comment = list(filter(None, self.comment))
+        self.comment.remove_empty()
 
         if not found and eval:
             self.comment.append(eval)
@@ -511,7 +575,7 @@ def arrows(self) -> List[chess.svg.Arrow]:
         Returns a list of :class:`arrows <chess.svg.Arrow>`.
         """
         arrows = []
-        for match in ARROWS_REGEX.finditer(" ".join(self.comment)):
+        for match in ARROWS_REGEX.finditer(self.comment.join(" ")):
             for group in match.group("arrows").split(","):
                 arrows.append(chess.svg.Arrow.from_pgn(group))
 
@@ -536,7 +600,7 @@ def set_arrows(self, arrows: Iterable[Union[chess.svg.Arrow, Tuple[Square, Squar
         for index in range(len(self.comment)):
             self.comment[index] = ARROWS_REGEX.sub(_condense_affix(""), self.comment[index])
 
-        self.comment = list(filter(None, self.comment))
+        self.comment.remove_empty()
 
         prefix = ""
         if csl:
@@ -555,7 +619,7 @@ def clock(self) -> Optional[float]:
         Returns the player's remaining time to the next time control after this
         move, in seconds.
         """
-        match = CLOCK_REGEX.search(" ".join(self.comment))
+        match = CLOCK_REGEX.search(self.comment.join(" "))
         if match is None:
             return None
         return int(match.group("hours")) * 3600 + int(match.group("minutes")) * 60 + float(match.group("seconds"))
@@ -580,7 +644,7 @@ def set_clock(self, seconds: Optional[float]) -> None:
             if found:
                 break
 
-        self.comment = list(filter(None, self.comment))
+        self.comment.remove_empty()
 
         if not found and clk:
             self.comment.append(clk)
@@ -593,7 +657,7 @@ def emt(self) -> Optional[float]:
         Returns the player's elapsed move time use for the comment of this
         move, in seconds.
         """
-        match = EMT_REGEX.search(" ".join(self.comment))
+        match = EMT_REGEX.search(self.comment.join(" "))
         if match is None:
             return None
         return int(match.group("hours")) * 3600 + int(match.group("minutes")) * 60 + float(match.group("seconds"))
@@ -618,7 +682,7 @@ def set_emt(self, seconds: Optional[float]) -> None:
             if found:
                 break
 
-        self.comment = list(filter(None, self.comment))
+        self.comment.remove_empty()
 
         if not found and emt:
             self.comment.append(emt)
@@ -677,7 +741,7 @@ class ChildNode(GameNode):
     move: chess.Move
     """The move leading to this node."""
 
-    starting_comment: list[str]
+    starting_comment: GameNodeComment
     """
     A comment for the start of a variation. Only nodes that
     actually start a variation (:func:`~chess.pgn.GameNode.starts_variation()`
@@ -698,7 +762,7 @@ def __init__(self, parent: GameNode, move: chess.Move, *, comment: Union[str, li
         self.parent.variations.append(self)
 
         self.nags.update(nags)
-        self.starting_comment = starting_comment if isinstance(starting_comment, list) else [starting_comment] if starting_comment else []
+        self.starting_comment = GameNodeComment(starting_comment)
 
     def board(self) -> chess.Board:
         stack: List[chess.Move] = []
@@ -1150,7 +1214,7 @@ def visit_board(self, board: chess.Board) -> None:
         """
         pass
 
-    def visit_comment(self, comment: list[str]) -> None:
+    def visit_comment(self, comment: GameNodeComment) -> None:
         """Called for each comment."""
         pass
 
@@ -1204,7 +1268,7 @@ def begin_game(self) -> None:
         self.game: GameT = self.Game()
 
         self.variation_stack: List[GameNode] = [self.game]
-        self.starting_comment: list[str] = []
+        self.starting_comment = GameNodeComment()
         self.in_variation = False
 
     def begin_headers(self) -> Headers:
@@ -1229,22 +1293,24 @@ def visit_result(self, result: str) -> None:
         if self.game.headers.get("Result", "*") == "*":
             self.game.headers["Result"] = result
 
-    def visit_comment(self, comment: list[str]) -> None:
+    def visit_comment(self, comment: GameNodeComment) -> None:
         if self.in_variation or (self.variation_stack[-1].parent is None and self.variation_stack[-1].is_end()):
             # Add as a comment for the current node if in the middle of
             # a variation. Add as a comment for the game if the comment
             # starts before any move.
             new_comment = self.variation_stack[-1].comment + comment
-            self.variation_stack[-1].comment = list(filter(None, new_comment))
+            new_comment.remove_empty()
+            self.variation_stack[-1].comment = new_comment
         else:
             # Otherwise, it is a starting comment.
             new_comment = self.starting_comment + comment
-            self.starting_comment = list(filter(None, new_comment))
+            new_comment.remove_empty()
+            self.starting_comment = new_comment
 
     def visit_move(self, board: chess.Board, move: chess.Move) -> None:
         self.variation_stack[-1] = self.variation_stack[-1].add_variation(move)
         self.variation_stack[-1].starting_comment = self.starting_comment
-        self.starting_comment = []
+        self.starting_comment = GameNodeComment()
         self.in_variation = True
 
     def handle_error(self, error: Exception) -> None:
@@ -1412,9 +1478,9 @@ def end_variation(self) -> None:
             self.write_token(") ")
             self.force_movenumber = True
 
-    def visit_comment(self, comment: list[str]) -> None:
+    def visit_comment(self, comment: GameNodeComment) -> None:
         if self.comments and (self.variations or not self.variation_depth):
-            self.write_token(" ".join("{ " + single_comment.replace("}", "").strip() + " }" for single_comment in comment) + " ")
+            self.write_token(comment.pgn_format() + " ")
             self.force_movenumber = True
 
     def visit_nag(self, nag: int) -> None:
@@ -1709,7 +1775,7 @@ def read_game(handle: TextIO, *, Visitor: Any = GameBuilder) -> Any:
                     line = line[close_index + 1:]
 
                 if not skip_variation_depth:
-                    visitor.visit_comment(["".join(comment_lines)])
+                    visitor.visit_comment(GameNodeComment("".join(comment_lines)))
 
                 # Continue with the current line.
                 fresh_line = False

test.py

@@ -2078,29 +2078,29 @@ class PgnTestCase(unittest.TestCase):
 
     def test_exporter(self):
         game = chess.pgn.Game()
-        game.comment = ["Test game:"]
+        game.comment.set("Test game:")
         game.headers["Result"] = "*"
         game.headers["VeryLongHeader"] = "This is a very long header, much wider than the 80 columns that PGNs are formatted with by default"
 
         e4 = game.add_variation(game.board().parse_san("e4"))
-        e4.comment = ["Scandinavian Defense:"]
+        e4.comment.set("Scandinavian Defense:")
 
         e4_d5 = e4.add_variation(e4.board().parse_san("d5"))
 
         e4_h5 = e4.add_variation(e4.board().parse_san("h5"))
         e4_h5.nags.add(chess.pgn.NAG_MISTAKE)
-        e4_h5.starting_comment = ["This"]
-        e4_h5.comment = ["is nonsense"]
+        e4_h5.starting_comment.set("This")
+        e4_h5.comment.set("is nonsense")
 
         e4_e5 = e4.add_variation(e4.board().parse_san("e5"))
         e4_e5_Qf3 = e4_e5.add_variation(e4_e5.board().parse_san("Qf3"))
         e4_e5_Qf3.nags.add(chess.pgn.NAG_MISTAKE)
 
         e4_c5 = e4.add_variation(e4.board().parse_san("c5"))
-        e4_c5.comment = ["Sicilian"]
+        e4_c5.comment.set("Sicilian")
 
         e4_d5_exd5 = e4_d5.add_main_variation(e4_d5.board().parse_san("exd5"))
-        e4_d5_exd5.comment = ["Best", "and the end of this example"]
+        e4_d5_exd5.comment.set(["Best", "and the end of this example"])
 
         # Test string exporter with various options.
         exporter = chess.pgn.StringExporter(headers=False, comments=False, variations=False)
@@ -2828,7 +2828,7 @@ def test_recursion(self):
 
     def test_annotations(self):
         game = chess.pgn.Game()
-        game.comment = ["foo [%bar] baz"]
+        game.comment = chess.pgn.GameNodeComment("foo [%bar] baz")
 
         self.assertTrue(game.clock() is None)
         clock = 12345
@@ -2879,7 +2879,7 @@ def test_eval(self):
 
     def test_float_emt(self):
         game = chess.pgn.Game()
-        game.comment = ["[%emt 0:00:01.234]"]
+        game.comment = chess.pgn.GameNodeComment("[%emt 0:00:01.234]")
         self.assertEqual(game.emt(), 1.234)
 
         game.set_emt(6.54321)
@@ -2892,7 +2892,7 @@ def test_float_emt(self):
 
     def test_float_clk(self):
         game = chess.pgn.Game()
-        game.comment = ["[%clk 0:00:01.234]"]
+        game.comment.set("[%clk 0:00:01.234]")
         self.assertEqual(game.clock(), 1.234)
 
         game.set_clock(6.54321)