From ec37db396c97267e48d2a7d7f5b29be17720f79f Mon Sep 17 00:00:00 2001
From: Niklas Fiekas <niklas.fiekas@backscattering.de>
Date: Thu, 28 Dec 2017 11:48:26 +0100
Subject: [PATCH] Add GameNode.accept_subgame

---
 chess/pgn.py | 25 +++++++++++++++++++++++++
 test.py      |  9 +++++++++
 2 files changed, 34 insertions(+)

diff --git a/chess/pgn.py b/chess/pgn.py
index 1bc2b8da3..455532c9f 100644
--- a/chess/pgn.py
+++ b/chess/pgn.py
@@ -353,6 +353,31 @@ def accept(self, visitor, _board=None):
         if _board is None:
             return visitor.result()
 
+    def accept_subgame(self, visitor):
+        """
+        Traverses headers and game nodes in PGN order, as if the game was
+        starting from this node. Returns the visitor result.
+        """
+        game = self.root()
+        visitor.begin_game()
+
+        dummy_game = Game.without_tag_roster()
+        dummy_game.setup(self.board())
+
+        visitor.begin_headers()
+        for tagname, tagvalue in game.headers.items():
+            if tagname not in dummy_game.headers:
+                visitor.visit_header(tagname, tagvalue)
+        for tagname, tagvalue in dummy_game.headers.items():
+            visitor.visit_header(tagname, tagvalue)
+        visitor.end_headers()
+
+        self.accept(visitor)
+
+        visitor.visit_result(game.headers.get("Result", "*"))
+        visitor.end_game()
+        return visitor.result()
+
     def __str__(self):
         return self.accept(StringExporter(columns=None))
 
diff --git a/test.py b/test.py
index 727b69efa..1c55497a6 100755
--- a/test.py
+++ b/test.py
@@ -2200,6 +2200,15 @@ def test_no_movetext(self):
 
         self.assertTrue(chess.pgn.read_game(pgn) is None)
 
+    def test_subgame(self):
+        pgn = StringIO("1. d4 d5 (1... Nf6 2. c4 (2. Nf3 g6 3. g3))")
+        game = chess.pgn.read_game(pgn)
+        node = game.variations[0].variations[1]
+        subgame = node.accept_subgame(chess.pgn.GameModelCreator())
+        self.assertEqual(subgame.headers["FEN"], "rnbqkb1r/pppppppp/5n2/8/3P4/8/PPP1PPPP/RNBQKBNR w KQkq - 1 2")
+        self.assertEqual(subgame.variations[0].move, chess.Move.from_uci("c2c4"))
+        self.assertEqual(subgame.variations[1].move, chess.Move.from_uci("g1f3"))
+
 
 class CraftyTestCase(unittest.TestCase):