From e7e65319c30ea5c2414200b44db29e693c610999 Mon Sep 17 00:00:00 2001 From: Niklas Fiekas Date: Fri, 16 Sep 2016 16:59:36 +0200 Subject: [PATCH] Add GameNode.main_line() --- README.rst | 5 +++++ chess/pgn.py | 8 +++++++- test.py | 15 +++++++++++---- 3 files changed, 23 insertions(+), 5 deletions(-) diff --git a/README.rst b/README.rst index 4b57d95d4..59780675d 100644 --- a/README.rst +++ b/README.rst @@ -271,6 +271,11 @@ Features >>> first_game.headers["Black"] 'Bordais' + >>> # Get the mainline as a list of moves. + >>> moves = first_game.main_line() + >>> first_game.board().variation_san(moves) + '1. e4 c5 2. c4 Nc6 3. Ne2 Nf6 4. Nbc3 Nb4 5. g3 Nd3#' + >>> # Iterate through the mainline of this embarrasingly short game. >>> node = first_game >>> while not node.is_end(): diff --git a/chess/pgn.py b/chess/pgn.py index d9f003ef9..a7964e384 100644 --- a/chess/pgn.py +++ b/chess/pgn.py @@ -246,6 +246,13 @@ def add_main_variation(self, move, comment=""): self.variations.insert(0, node) return node + def main_line(self): + """Yields the moves of the main line starting in this node.""" + node = self + while node.variations: + node = node.variations[0] + yield node.move + def add_line(self, moves, comment="", starting_comment="", nags=()): """ Creates a sequence of child nodes for the given list of moves. @@ -268,7 +275,6 @@ def add_line(self, moves, comment="", starting_comment="", nags=()): return node - def accept(self, visitor, _board=None): """ Traverse game nodes in PGN order using the given *visitor*. Returns diff --git a/test.py b/test.py index c57900649..98ff01d85 100755 --- a/test.py +++ b/test.py @@ -1966,10 +1966,9 @@ def test_add_line(self): game = chess.pgn.Game() game.add_variation(chess.Move.from_uci("e2e4")) - tail = game.add_line([ - chess.Move.from_uci("g1f3"), - chess.Move.from_uci("d7d5") - ], starting_comment="start", comment="end", nags=(17, 42)) + moves = [chess.Move.from_uci("g1f3"), chess.Move.from_uci("d7d5")] + + tail = game.add_line(moves, starting_comment="start", comment="end", nags=(17, 42)) self.assertEqual(tail.parent.move, chess.Move.from_uci("g1f3")) self.assertEqual(tail.parent.starting_comment, "start") @@ -1980,6 +1979,14 @@ def test_add_line(self): self.assertEqual(tail.comment, "end") self.assertTrue(42 in tail.nags) + def test_main_line(self): + moves = [chess.Move.from_uci(uci) for uci in ["d2d3", "g8f6", "e2e4"]] + + game = chess.pgn.Game() + game.add_line(moves) + + self.assertEqual(list(game.main_line()), moves) + class StockfishTestCase(unittest.TestCase):